# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
import numpy as np
from qutip.operators import sigmax, sigmay, sigmaz, identity
from qutip.tensor import tensor
from qutip.qip.circuit import QubitCircuit
from qutip.qip.device.processor import Processor
from qutip.qip.device.modelprocessor import ModelProcessor, GateDecomposer

__all__ = ['SpinChain', 'LinearSpinChain', 'CircularSpinChain',
           'SpinChainGateDecomposer']

class SpinChain(ModelProcessor):
    """
    The processor based on the physical implementation of
    a spin chain qubits system.
    The available Hamiltonian of the system is predefined.
    The processor can simulate the evolution under the given
    control pulses either numerically or analytically.
    It is a base class and should not be used directly, please
    refer the the subclasses :class:`qutip.qip.LinearSpinChain` and
    :class:`qutip.qip.CircularSpinChain`.

    Parameters
    ----------
    N: int
        The number of qubits in the system.

    correct_global_phase: float
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    sx: int or list
        The delta for each of the qubits in the system.

    sz: int or list
        The epsilon for each of the qubits in the system.

    sxsy: int or list
        The interaction strength for each of the qubit pair in the system.

    t1: list or float
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size ``N`` or a float for all qubits.

    t2: list of float
        Characterize the decoherence of dephasing for
        each qubit. A list of size ``N`` or a float for all qubits.

    Attributes
    ----------
    N: int
        The number of component systems.

    ctrls: list
        A list of the control Hamiltonians driving the evolution.

    tlist: array_like
        A NumPy array specifies the time of each coefficient.

    coeffs: array_like
        A 2d NumPy array of the shape, the length is dependent on the
        spline type

    t1: list
        Characterize the decoherence of amplitude damping for
        each qubit.

    t2: list
        Characterize the decoherence of dephasing for
        each qubit.

    noise: :class:`qutip.qip.Noise`, optional
        A list of noise objects. They will be processed when creating the
        noisy :class:`qutip.QobjEvo` from the processor or run the simulation.

    dims: list
        The dimension of each component system.
        Default is dim=[2,2,2,...,2]

    spline_kind: str
        Type of the coefficient interpolation.
        Note that they have different requirement for the length of ``coeffs``.

        -"step_func":
        The coefficient will be treated as a step function.
        E.g. ``tlist=[0,1,2]`` and ``coeffs=[3,2]``, means that the coefficient
        is 3 in t=[0,1) and 2 in t=[2,3). It requires
        ``coeffs.shape[1]=len(tlist)-1`` or ``coeffs.shape[1]=len(tlist)``, but
        in the second case the last element has no effect.

        -"cubic": Use cubic interpolation for the coefficient. It requires
        ``coeffs.shape[1]=len(tlist)``

    sx: list
        The delta for each of the qubits in the system.

    sz: list
        The epsilon for each of the qubits in the system.

    sxsy: list
        The interaction strength for each of the qubit pair in the system.

    sx_ops: list
        A list of sigmax Hamiltonians for each qubit.

    sz_ops: list
        A list of sigmaz Hamiltonians for each qubit.

    sxsy_ops: list
        A list of tensor(sigmax, sigmay)
        interacting Hamiltonians for each qubit.

    sx_u: array_like
        Pulse matrix for sigmax Hamiltonians.

    sz_u: array_like
        Pulse matrix for sigmaz Hamiltonians.

    sxsy_u: array_like
        Pulse matrix for tensor(sigmax, sigmay) interacting Hamiltonians.
    """
    def __init__(self, N, correct_global_phase,
                 sx, sz, sxsy, t1, t2):
        super(SpinChain, self).__init__(
            N, correct_global_phase=correct_global_phase, t1=t1, t2=t2)
        self.correct_global_phase = correct_global_phase
        self.spline_kind = "step_func"
        # params and ops are set in the submethods

    def set_up_ops(self, N):
        """
        Generate the Hamiltonians for the spinchain model and save them in the
        attribute `ctrls`.

        Parameters
        ----------
        N: int
            The number of qubits in the system.
        """
        # sx_ops
        self.ctrls += [tensor([sigmax() if m == n else identity(2)
                               for n in range(N)])
                       for m in range(N)]
        # sz_ops
        self.ctrls += [tensor([sigmaz() if m == n else identity(2)
                               for n in range(N)])
                       for m in range(N)]
        # sxsy_ops
        for n in range(N - 1):
            x = [identity(2)] * N
            x[n] = x[n + 1] = sigmax()
            y = [identity(2)] * N
            y[n] = y[n + 1] = sigmay()
            self.ctrls.append(tensor(x) + tensor(y))

    def set_up_params(self, sx, sz):
        """
        Save the parameters in the attribute `params` and check the validity.

        Parameters
        ----------
        sx: float or list
            The coefficient of sigmax in the model

        sz: flaot or list
            The coefficient of sigmaz in the model

        Notes
        -----
        The coefficient of sxsy is defined in the submethods.
        All parameters will be multiplied by 2*pi for simplicity
        """
        sx_para = super(SpinChain, self)._para_list(sx, self.N)
        self._paras["sx"] = sx_para
        sz_para = super(SpinChain, self)._para_list(sz, self.N)
        self._paras["sz"] = sz_para

    @property
    def sx_ops(self):
        return self.ctrls[: self.N]

    @property
    def sz_ops(self):
        return self.ctrls[self.N: 2*self.N]

    @property
    def sxsy_ops(self):
        return self.ctrls[2*self.N:]

    @property
    def sx_u(self):
        return self.coeffs[: self.N]

    @property
    def sz_u(self):
        return self.coeffs[self.N: 2*self.N]

    @property
    def sxsy_u(self):
        return self.coeffs[2*self.N:]

    def load_circuit(self, qc, setup):
        """
        Decompose a :class:`qutip.QubitCircuit` in to the control
        amplitude generating the corresponding evolution.

        Parameters
        ----------
        qc: :class:`qutip.QubitCircuit`
            Takes the quantum circuit to be implemented.

        setup: string
            "linear" or "circular" for two sub-calsses.

        Returns
        -------
        tlist: array_like
            A NumPy array specifies the time of each coefficient

        coeffs: array_like
            A 2d NumPy array of the shape (len(ctrls), len(tlist)). Each
            row corresponds to the control pulse sequence for
            one Hamiltonian.
        """
        gates = self.optimize_circuit(qc).gates

        dec = SpinChainGateDecomposer(
            self.N, self._paras, setup=setup,
            global_phase=0., num_ops=len(self.ctrls))
        self.tlist, self.coeffs, self.global_phase = dec.decompose(gates)

        return self.tlist, self.coeffs

    def adjacent_gates(self, qc, setup="linear"):
        """
        Method to resolve 2 qubit gates with non-adjacent control/s or target/s
        in terms of gates with adjacent interactions for linear/circular spin
        chain system.

        Parameters
        ----------
        qc: :class:`qutip.QubitCircuit`
            The circular spin chain circuit to be resolved

        setup: Boolean
            Linear of Circular spin chain setup

        Returns
        -------
        qc: :class:`qutip.QubitCircuit`
            Returns QubitCircuit of resolved gates for the qubit circuit in the
            desired basis.
        """
        qc_t = QubitCircuit(qc.N, qc.reverse_states)
        swap_gates = ["SWAP", "ISWAP", "SQRTISWAP", "SQRTSWAP", "BERKELEY",
                      "SWAPalpha"]
        N = qc.N

        for gate in qc.gates:
            if gate.name == "CNOT" or gate.name == "CSIGN":
                start = min([gate.targets[0], gate.controls[0]])
                end = max([gate.targets[0], gate.controls[0]])

                if (setup == "linear" or
                        (setup == "circular" and (end - start) <= N // 2)):
                    i = start
                    while i < end:
                        if (start + end - i - i == 1 and
                                (end - start + 1) % 2 == 0):
                            # Apply required gate if control and target are
                            # adjacent to each other, provided |control-target|
                            # is even.
                            if end == gate.controls[0]:
                                qc_t.add_gate(gate.name, targets=[i],
                                              controls=[i + 1])
                            else:
                                qc_t.add_gate(gate.name, targets=[i + 1],
                                              controls=[i])

                        elif (start + end - i - i == 2 and
                              (end - start + 1) % 2 == 1):
                            # Apply a swap between i and its adjacent gate,
                            # then the required gate if and then another swap
                            # if control and target have one qubit between
                            # them, provided |control-target| is odd.
                            qc_t.add_gate("SWAP", targets=[i, i + 1])
                            if end == gate.controls[0]:
                                qc_t.add_gate(gate.name, targets=[i + 1],
                                              controls=[i + 2])
                            else:
                                qc_t.add_gate(gate.name, targets=[i + 2],
                                              controls=[i + 1])
                            qc_t.add_gate("SWAP", [i, i + 1])
                            i += 1

                        else:
                            # Swap the target/s and/or control with their
                            # adjacent qubit to bring them closer.
                            qc_t.add_gate("SWAP", [i, i + 1])
                            qc_t.add_gate("SWAP", [start + end - i - 1,
                                                   start + end - i])
                        i += 1

                elif (end - start) < N - 1:
                    """
                    If the resolving has to go backwards, the path is first
                    mapped to a separate circuit and then copied back to the
                    original circuit.
                    """

                    temp = QubitCircuit(N - end + start)
                    i = 0
                    while i < (N - end + start):

                        if (N + start - end - i - i == 1 and
                                (N - end + start + 1) % 2 == 0):
                            if end == gate.controls[0]:
                                temp.add_gate(gate.name, targets=[i],
                                              controls=[i + 1])
                            else:
                                temp.add_gate(gate.name, targets=[i + 1],
                                              controls=[i])

                        elif (N + start - end - i - i == 2 and
                              (N - end + start + 1) % 2 == 1):
                            temp.add_gate("SWAP", targets=[i, i + 1])
                            if end == gate.controls[0]:
                                temp.add_gate(gate.name, targets=[i + 2],
                                              controls=[i + 1])
                            else:
                                temp.add_gate(gate.name, targets=[i + 1],
                                              controls=[i + 2])
                            temp.add_gate("SWAP", [i, i + 1])
                            i += 1

                        else:
                            temp.add_gate("SWAP", [i, i + 1])
                            temp.add_gate("SWAP",
                                          [N + start - end - i - 1,
                                           N + start - end - i])
                        i += 1

                    j = 0
                    for gate in temp.gates:
                        if (j < N - end - 2):
                            if gate.name in ["CNOT", "CSIGN"]:
                                qc_t.add_gate(gate.name, end + gate.targets[0],
                                              end + gate.controls[0])
                            else:
                                qc_t.add_gate(gate.name,
                                              [end + gate.targets[0],
                                               end + gate.targets[1]])
                        elif (j == N - end - 2):
                            if gate.name in ["CNOT", "CSIGN"]:
                                qc_t.add_gate(gate.name, end + gate.targets[0],
                                              (end + gate.controls[0]) % N)
                            else:
                                qc_t.add_gate(gate.name,
                                              [end + gate.targets[0],
                                               (end + gate.targets[1]) % N])
                        else:
                            if gate.name in ["CNOT", "CSIGN"]:
                                qc_t.add_gate(gate.name,
                                              (end + gate.targets[0]) % N,
                                              (end + gate.controls[0]) % N)
                            else:
                                qc_t.add_gate(gate.name,
                                              [(end + gate.targets[0]) % N,
                                               (end + gate.targets[1]) % N])
                        j = j + 1

                elif (end - start) == N - 1:
                    qc_t.add_gate(gate.name, gate.targets, gate.controls)

            elif gate.name in swap_gates:
                start = min([gate.targets[0], gate.targets[1]])
                end = max([gate.targets[0], gate.targets[1]])

                if (setup == "linear" or
                        (setup == "circular" and (end - start) <= N // 2)):
                    i = start
                    while i < end:
                        if (start + end - i - i == 1 and
                                (end - start + 1) % 2 == 0):
                            qc_t.add_gate(gate.name, [i, i + 1])
                        elif ((start + end - i - i) == 2 and
                              (end - start + 1) % 2 == 1):
                            qc_t.add_gate("SWAP", [i, i + 1])
                            qc_t.add_gate(gate.name, [i + 1, i + 2])
                            qc_t.add_gate("SWAP", [i, i + 1])
                            i += 1
                        else:
                            qc_t.add_gate("SWAP", [i, i + 1])
                            qc_t.add_gate("SWAP", [start + end - i - 1,
                                                   start + end - i])
                        i += 1

                else:
                    temp = QubitCircuit(N - end + start)
                    i = 0
                    while i < (N - end + start):

                        if (N + start - end - i - i == 1 and
                                (N - end + start + 1) % 2 == 0):
                            temp.add_gate(gate.name, [i, i + 1])

                        elif (N + start - end - i - i == 2 and
                              (N - end + start + 1) % 2 == 1):
                            temp.add_gate("SWAP", [i, i + 1])
                            temp.add_gate(gate.name, [i + 1, i + 2])
                            temp.add_gate("SWAP", [i, i + 1])
                            i += 1

                        else:
                            temp.add_gate("SWAP", [i, i + 1])
                            temp.add_gate("SWAP", [N + start - end - i - 1,
                                                   N + start - end - i])
                        i += 1

                    j = 0
                    for gate in temp.gates:
                        if(j < N - end - 2):
                            qc_t.add_gate(gate.name, [end + gate.targets[0],
                                                      end + gate.targets[1]])
                        elif(j == N - end - 2):
                            qc_t.add_gate(gate.name,
                                          [end + gate.targets[0],
                                           (end + gate.targets[1]) % N])
                        else:
                            qc_t.add_gate(gate.name,
                                          [(end + gate.targets[0]) % N,
                                           (end + gate.targets[1]) % N])
                        j = j + 1

            else:
                qc_t.add_gate(gate.name, gate.targets, gate.controls,
                              gate.arg_value, gate.arg_label)

        return qc_t

    def eliminate_auxillary_modes(self, U):
        return U

    def optimize_circuit(self, qc):
        """
        Take a quantum circuit/algorithm and convert it into the
        optimal form/basis for the desired physical system.

        Parameters
        ----------
        qc: :class:`qutip.QubitCircuit`
            Takes the quantum circuit to be implemented.

        Returns
        -------
        qc: :class:`qutip.QubitCircuit`
            The circuit representation with elementary gates
            that can be implemented in this model.
        """
        self.qc0 = qc
        self.qc1 = self.adjacent_gates(self.qc0)
        self.qc2 = self.qc1.resolve_gates(
            basis=["SQRTISWAP", "ISWAP", "RX", "RZ"])
        return self.qc2

class LinearSpinChain(SpinChain):
    """
    A processor based on the physical implementation of
    a linear spin chain qubits system.
    The available Hamiltonian of the system is predefined.
    The processor can simulate the evolution under the given
    control pulses either numerically or analytically.

    Parameters
    ----------
    N: int
        The number of qubits in the system.

    correct_global_phase: float
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    sx: int or list
        The delta for each of the qubits in the system.

    sz: int or list
        The epsilon for each of the qubits in the system.

    sxsy: int or list
        The interaction strength for each of the qubit pair in the system.

    t1: list or float, optional
        Characterize the decoherence of amplitude damping for
        each qubit.

    t2: list of float, optional
        Characterize the decoherence of dephasing for
        each qubit.
    """
    def __init__(self, N, correct_global_phase=True,
                 sx=0.25, sz=1.0, sxsy=0.1, t1=None, t2=None):

        super(LinearSpinChain, self).__init__(
            N, correct_global_phase=correct_global_phase,
            sx=sx, sz=sz, sxsy=sxsy, t1=t1, t2=t2)
        self.set_up_params(sx=sx, sz=sz, sxsy=sxsy)
        self.set_up_ops(N)

    def set_up_ops(self, N):
        super(LinearSpinChain, self).set_up_ops(N)

    def set_up_params(self, sx, sz, sxsy):
        # Doc same as in the parent class
        super(LinearSpinChain, self).set_up_params(sx, sz)
        sxsy_para = self._para_list(sxsy, self.N-1)
        self._paras["sxsy"] = sxsy_para

    @property
    def sxsy_ops(self):
        return self.ctrls[2*self.N: 3*self.N-1]

    @property
    def sxsy_u(self):
        return self.coeffs[2*self.N: 3*self.N-1]

    def load_circuit(self, qc):
        return super(LinearSpinChain, self).load_circuit(qc, "linear")

    def get_ops_labels(self):
        return ([r"$\sigma_x^%d$" % n for n in range(self.N)] +
                [r"$\sigma_z^%d$" % n for n in range(self.N)] +
                [r"$\sigma_x^%d\sigma_x^{%d} + \sigma_y^%d\sigma_y^{%d}$"
                 % (n, n, n + 1, n + 1) for n in range(self.N - 1)])

    def adjacent_gates(self, qc):
        return super(LinearSpinChain, self).adjacent_gates(qc, "linear")

class CircularSpinChain(SpinChain):
    """
    A processor based on the physical implementation of
    a circular spin chain qubits system.
    The available Hamiltonian of the system is predefined.
    The processor can simulate the evolution under the given
    control pulses either numerically or analytically.

    Parameters
    ----------
    N: int
        The number of qubits in the system.

    correct_global_phase: float
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    sx: int or list
        The delta for each of the qubits in the system.

    sz: int or list
        The epsilon for each of the qubits in the system.

    sxsy: int or list
        The interaction strength for each of the qubit pair in the system.

    t1: list or float, optional
        Characterize the decoherence of amplitude damping for
        each qubit.

    t2: list of float, optional
        Characterize the decoherence of dephasing for
        each qubit.
    """
    def __init__(self, N, correct_global_phase=True,
                 sx=0.25, sz=1.0, sxsy=0.1, t1=None, t2=None):

        super(CircularSpinChain, self).__init__(
            N, correct_global_phase=correct_global_phase,
            sx=sx, sz=sz, sxsy=sxsy, t1=t1, t2=t2)
        self.set_up_params(sx=sx, sz=sz, sxsy=sxsy)
        self.set_up_ops(N)

    def set_up_ops(self, N):
        super(CircularSpinChain, self).set_up_ops(N)
        x = [identity(2)] * N
        x[0] = x[N - 1] = sigmax()
        y = [identity(2)] * N
        y[0] = y[N - 1] = sigmay()
        self.ctrls.append(tensor(x) + tensor(y))

    def set_up_params(self, sx, sz, sxsy):
        # Doc same as in the parent class
        super(CircularSpinChain, self).set_up_params(sx, sz)
        sxsy_para = self._para_list(sxsy, self.N)
        self._paras["sxsy"] = sxsy_para

    @property
    def sxsy_ops(self):
        return self.ctrls[2*self.N: 3*self.N]

    @property
    def sxsy_u(self):
        return self.coeffs[2*self.N: 3*self.N]

    def load_circuit(self, qc):
        return super(CircularSpinChain, self).load_circuit(qc, "circular")

    def get_ops_labels(self):
        return ([r"$\sigma_x^%d$" % n for n in range(self.N)] +
                [r"$\sigma_z^%d$" % n for n in range(self.N)] +
                [r"$\sigma_x^%d\sigma_x^{%d} + \sigma_y^%d\sigma_y^{%d}$"
                 % (n, n, (n + 1) % self.N, (n + 1) % self.N)
                 for n in range(self.N)])

    def adjacent_gates(self, qc):
        return super(CircularSpinChain, self).adjacent_gates(qc, "circular")

class SpinChainGateDecomposer(GateDecomposer):
    """
    Decompose a :class:`qutip.QubitCircuit` into
    the pulse sequence for the processor.

    Parameters
    ----------
    N: int
        The number of qubits in the system.

    params: dict
        A Python dictionary contains the name and the value of the parameters,
        such as laser frequency, detuning etc.

    setup: string
        "linear" or "circular" for two sub-calsses.

    global_phase: bool
        Record of the global phase change and will be returned.

    num_ops: int
        Number of Hamiltonians in the processor.

    Attributes
    ----------
    N: int
        The number of the component systems.

    params: dict
        A Python dictionary contains the name and the value of the parameters,
        such as laser frequency, detuning etc.

    num_ops: int
        Number of control Hamiltonians in the processor.

    gate_decomps: dict
        The Python dictionary in the form of {gate_name: decompose_function}.
        It saves the decomposition scheme for each gate.

    setup: string
        "linear" or "circular" for two sub-calsses.

    global_phase: bool
        Record of the global phase change and will be returned.
    """
    def __init__(self, N, params, setup, global_phase, num_ops):
        super(SpinChainGateDecomposer, self).__init__(
            N=N, params=params, num_ops=num_ops)
        self.gate_decomps = {"ISWAP": self.iswap_dec,
                             "SQRTISWAP": self.sqrtiswap_dec,
                             "RZ": self.rz_dec,
                             "RX": self.rx_dec,
                             "GLOBALPHASE": self.globalphase_dec
                             }
        self.N = N
        self._sx_ind = list(range(0, N))
        self._sz_ind = list(range(N, 2*N))
        if setup == "circular":
            self._sxsy_ind = list(range(2*N, 3*N))
        elif setup == "linear":
            self._sxsy_ind = list(range(2*N, 3*N-1))
        self.global_phase = global_phase

    def decompose(self, gates):
        tlist, coeffs = super(SpinChainGateDecomposer, self).decompose(gates)
        return tlist, coeffs, self.global_phase

    def rz_dec(self, gate):
        """
        Decomposer for the RZ gate
        """
        pulse = np.zeros(self.num_ops)
        q_ind = gate.targets[0]
        g = self.params["sz"][q_ind]
        pulse[self._sz_ind[q_ind]] = np.sign(gate.arg_value) * g
        t = abs(gate.arg_value) / (2 * g)
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

    def rx_dec(self, gate):
        """
        Decomposer for the RX gate
        """
        pulse = np.zeros(self.num_ops)
        q_ind = gate.targets[0]
        g = self.params["sx"][q_ind]
        pulse[self._sx_ind[q_ind]] = np.sign(gate.arg_value) * g
        t = abs(gate.arg_value) / (2 * g)
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

    def iswap_dec(self, gate):
        """
        Decomposer for the ISWAP gate
        """
        pulse = np.zeros(self.num_ops)
        q1, q2 = min(gate.targets), max(gate.targets)
        g = self.params["sxsy"][q1]
        if q1 == 0 and q2 == self.N - 1:
            pulse[self._sxsy_ind[self.N - 1]] = -g
        else:
            pulse[self._sxsy_ind[q1]] = -g
        t = np.pi / (4 * g)
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

    def sqrtiswap_dec(self, gate):
        """
        Decomposer for the SQRTISWAP gate
        """
        pulse = np.zeros(self.num_ops)
        q1, q2 = min(gate.targets), max(gate.targets)
        g = self.params["sxsy"][q1]
        if q1 == 0 and q2 == self.N - 1:
            pulse[self._sxsy_ind[self.N - 1]] = -g
        else:
            pulse[self._sxsy_ind[q1]] = -g
        t = np.pi / (8 * g)
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

    def globalphase_dec(self, gate):
        """
        Decomposer for the GLOBALPHASE gate
        """
        self.global_phase += gate.arg_value


import re
import numpy as np
import pandas as pd
import pickle
import random
import datetime
import argparse
import warnings
from collections import defaultdict
from Bio import SeqIO

def read_fasta(input_fasta):
    """ Read in human genome reference FASTA"""
    genome = {}
    for seq_record in SeqIO.parse(input_fasta, "fasta"):
        genome[seq_record.id] = seq_record.upper().seq.tomutable()
    ## remove all 'useless' chromosomes, i.e. must match chrX, chrY or "^[a-z]{3}\d{1,2}$"
    genome = {k: v for k, v in genome.items()
                    if re.match('^[a-z]{3}\d{1,2}$', k, re.IGNORECASE) or k in ['chrX', 'chrY', 'chrx', 'chry'] or k in ['X', 'Y', 'x', 'y']}
    return genome

def get_reference_bases(my_list):
	""" Convert REF base to ALT base such that REF != ALT"""
    bases = set('ATCG')
    legend = {base:list(bases - set(base)) for base in bases}  ## key-value legend, `{'C': ['T', 'A', 'G'], ...`

    new_list = [random.choice(legend[base]) for base in my_list]
    return new_list[0]

def retrieve_snp_positions(column):
    """For SNPs, retrieve 'number' positions between 'start' and 'end' """
    if int(column.number <= abs(int(column.end + 1) - int(column.start)):
        choice = np.random.choice(np.arange(column.start, column.end+1), column.number).tolist()
    else:
        raise ValueError('Input variant *csv is formatted incorrectly. Number of SNPs exceeds area to place them. Check input *csv columns "number", "length" or "position" ')

    return choice

def retrieve_del_positions(column):
    """For DELs, retrieve bases from reference length 'length', /times 'number' """
    ## retrieve random starts, such that these ared placed within 'length'+1
    arr = column.end - column.start - (column.number - 1) * column.length
    if arr <= 0:
         raise ValueError('Input variant *csv is formatted incorrectly. Deletions cannot fit in the interval (start, end). Either "number" &/or "length" are too large, or the interval too short.')
    elif arr > column.number:
        raise ValueError('Input variant *csv is formatted incorrectly. There are too many deletions requested within the interval (start, end). Either "number" &/or "length" are too large, or the interval too short.')
    choice = np.random.choice(column.end - column.start - (column.number - 1) * column.length, column.number, replace=False) ## draw N==number amount of unique choices
    idx = np.argsort(choice)  ## return indices that sort the above
    ## now add the "spacing" from 'length' after the draws, to make sure deletions do not clash with each other
    choice[idx] += np.arange(column.start, column.start + column.length * column.number, column.length)
    return choice

def read_variant_inputs(variants_csv, annotation_file):
    """ Read csv of variants for VCF, outputs dictionary of all positions.
        If only gene names provided, randomly draw N variants between gene_start and gene_end"""
    variants = pd.read_csv(variants_csv)
    variants.columns = [x.lower() for x in variants.columns]  ## format into all lowercase
    ## check both mandatory column pairs exist
    col_pairs = [['gene', 'number'], ['chr', 'position']]
    for pair in col_pairs:
        if any(col in variants for col in pair) and all(col in variants for col in pair):
            break
        else:
            raise ValueError('Input variant *csv is formatted incorrectly. Mandatory columns must be included. See README for formatting requirements.')
    ## if "variant_type" left empty, fill with "SNP" by default
    if 'variant_type' not in variants:
    	variants['variant_type'] = 'SNP'
    else:
    	variants['variant_type'].fillna('SNP', inplace=True)
    ## check elements in "variant_type" are valid
    if not set(['SNP', 'DEL', 'INS', 'DUP', 'INV', 'CNV', 'DUP:TANDEM', 'DEL:ME', 'INS:ME']).issubset(variants['variant_type']):
        raise ValueError('Input variant *csv is formatted incorrectly. The "variant_type" column must contain correctly labeled variants. See README for formatting requirements.')
    if set(['DEL:ME']).issubset(variants['variant_type']):
        warnings.warn("Warning: DEL:ME denotes a deletion of mobile element relative to the reference. This needs to be implemented by the user.")
    if set(['INS:ME']).issubset(variants['variant_type']):
        warnings.warn("Warning: INS:ME denotes an insertion of mobile element relative to the reference. This needs to be implemented by the user.")
    if set(['DUP']).issubset(variants['variant_type']):
        warnings.warn("Warning: DUP is implemented as an insertion of a nearby sequence of length=length.")

    ## if "position" exists, fill NA in "length" with 1
    if 'position' in variants.columns:
        if 'length' not in variants:   ## if "length" left empty, fill with 1 by default
            variants['length'] = 1
        else:
    	    variants['length'].fillna(1, inplace=True)

    if variants['position'].max() > 2.49e8:
        raise ValueError('Input variant *csv is formatted incorrectly. Value in `position` column greater than 2.49e8 bp, longer than the largest chromosome in the human reference. See README for formatting requirements.')

    ## check for required `length` column for insertions and deletions
    if set(['DEL', 'INS', 'DUP', 'INV', 'CNV', 'DUP:TANDEM', 'DEL:ME', 'INS:ME']).issubset(variants['variant_type']):
    	if 'length' not in variants:
    		variants['length'] = 1
        else:
            variants['length'].fillna(1, inplace=True)
    ## end = position if variant == SNP, else end = start + length
    variants['end'] = np.where(variants['variant_type']=='SNP', df['position'], df['start'] + df['length'])
    ### Adopt consistent name convention
    if 'gene' in variants.columns:
        variants = variants.rename(columns={'gene': 'gene_name'})   ## or re-write the annotation convention for 'gene' instead of 'gene_name'
    elif 'position' in variants.columns:
    	variants = variants.rename(columns={'chr': 'seqnames', 'position':'start'})  ## or re-write the annotation convertion for `chromosome` instead of `seqnames`
    ## subset annotation file to include chrom, start, end, and gene name at that position
    annotation = annotation[['seqname', 'start', 'end', 'gene_name']]
    ## If gene, then find gene name and get N numbers between (start, end)
    if 'gene_name' in variants.columns:
        match = variants.merge(annotation, on='gene_name', how='left')
        ## randomly get N positions between (start, end)
        #CODE     match['new_variant_positions'] = match.query('variant_type == "SNP"').apply(retrieve_snp_positions, axis=1)
        ## convert df to dictionary
        #CODE     variants_dict = match.groupby('seqname')['new_variant_positions'].apply(sum).to_dict()
    else:
    	match = variants  ## if user didn't provide gene_names, the annotation isn't useful/needed
    	### must convert to variant dictionary
    	### every `chr` must become the key, and every position a value in the key
    	#CODE     variants_dict = match.groupby('start').agg(lambda x: x.values.tolist()).T
    return match

def generate_snps_dict(match_df):
    """Generate SNPs from match DataFrame, output variants_dict"""
    match_df['new_variant_positions'] = match_df.query('variant_type == "SNP"').apply(retrieve_snp_positions, axis=1)
    snps_only = match_df[match_df.variant_type == 'SNP']
    if 'gene_name' in snps_only.columns:
        variants_dict = snps_only.groupby('seqname')['new_variant_positions'].apply(sum).to_dict()
    else:
        variants_dict = match_df.groupby('start').agg(lambda x: x.values.tolist()).T
    return variants_dicts

def generate_del_dict(match_df):
    """Generate DELs from match DataFrame, output variants_dict"""

def create_ref_alt_dict(var_dict, reference):
    """ Input dictionary of variant positions and return dictionary of positions w/ REF and ALT"""
    sorted_dict = {key: sorted(val) for key, val in var_dict.items()} ## sort dict values, which are lists
    ## del var_dict
    d = defaultdict(list)
    for key, val in sorted_dict.items():
    	for position in val:
    		ref = reference[key][position]
    		alt = get_reference_bases(ref)
    		d[key][position].append(ref)    ## make `position` the key, REF first, ALT second
            d[key][position].append(alt)

    del sorted_dict
    return d

def write_vcf(ref_alt_dict, outputVCF):
    """Write VCF from dictionary of positions/REF/variants"""
    df = pd.DataFrame.from_records([[i, j] + ref_alt_dict[i][j] for i in ref_alt_dict for j in ref_alt_dict[i]])
    df.columns = ['CHROM', 'POS', 'REF', 'ALT']

    df['ID'] = 'SNP'  ## must change for more variants
    df['QUAL'] = 40
    df['FILTER'] = 'PASS'
    df['INFO'] = 'simulated'

    df = df[['CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER', 'INFO']]

	now = datetime.datetime.now()

    header = '''##fileformat=VCFv4.1
    ##fileDate=''' + str(now.year) + str(now.day) + str(now.month) + '''
    ##source=myImputationProgramV3.1
    ##reference=file:///seq/references/
    #CHROM POS ID REF ALT QUAL FILTER INFO
    '''

    with open(outputVCF, 'w') as vcf:
        vcf.write(header)

    df.to_csv(outputVCF, sep='\t', mode='a', index=False)  ##  'a' appends to the end of the header above

def set_random_seed(seed_number):
    """Set numpy.random.seed for reproducible variant placements"""
    try:
        seed_input = int(seed_number)
    except:
        raise TypeError("Input must be a number!")
    return np.random.seed(seed=seed_input)

def main(args):
    ## check if seed set
    if args['seed'] is not None:
        set_random_seed(args['seed'])
    ## read in annotation of variants
    if args['reference']=='hg19' or args['reference']=='GRCh37':
        if args['annotation']=='GENCODE':
            annotation_gencode = pd.read_csv('../data/GENCODE_hg19_annotation_unique_gene_names.csv')
            vars = read_variant_inputs(args['input_variants'], annotation_gencode)
        elif args['annotation']=='Ensembl':
            annotation_ensembl = pd.read_csv('../data/ENSEMBL_Homo_sapiens.GRCh37.75_unique_gene_names.csv')
            vars = read_variant_inputs(args['input_variants'], annotation_ensembl)
        else:
            raise NotImplementedError('Error: "annotation" invalid. Only GENCODE and Ensembl annotations implemented.')
    elif args['reference']=='hg38' or args['reference']=='GRCh38':
        if args['annotation']=='GENCODE':
            annotation_gencode = pd.read_csv('../data/GENCODE_annotation_unique_gene_names.csv')
            vars = read_variant_inputs(args['input_variants'], annotation_gencode)
        elif args['annotation']=='Ensembl':
            annotation_ensembl = pd.read_csv('../data/ENSEMBL_Homo_sapiens.GRCh37.75_unique_gene_names.csv')
            vars = read_variant_inputs(args['input_variants'], annotation_ensembl)
        else:
            raise NotImplementedError('Error: "annotation" invalid. Only GENCODE and Ensembl annotations implemented.')
    else:
        raise NotImplementedError('Error: "reference" invalid. Only implemented with human genome reference hg19/GRCh37 or hg38/GRCh38.')
    ref_alts = create_ref_alt_dict(vars, args['input_fasta'])
    write_vcf(ref_alts, args['output_vcf'])

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Simulate VCF of somatic variants')
    parser.add_argument('-r','--reference',
                        choices=['hg19', 'GRCh37', 'hg38', 'GRCh38'],
                        type=str.lower,
                        default='hg19',
                        help='reference assembly; default hg19')
    parser.add_argument('-a','--annotation',
                        choices=['GENCODE', 'Ensembl'],
                        type=str.lower,
                        default='GENCODE',
                        help='annotation flag corresonding to input_fasta')
    parser.add_argument('-i','--input_fasta',
                        default='../data/subsampled_hg38.fa',
                        help='file path for the input (default genome) fasta')
    parser.add_argument('-v', '--input_variants',
                        default='../data/somatic_variants.csv',
                        help='file path for the input *csv of variants')
    parser.add_argument('-o', '--output_vcf',
                        default = 'outputs/somatic.VCF',
                        help='file path for the output VCF of somatic variants')
    parser.add_argument('-s', '--seed',
                        help='set seed for numpy.random.seed()')
    args = vars(parser.parse_args())
    main(args)


#!/usr/bin/env python3
import redis
import json
import os
import numpy as np
redis_connection = redis.Redis(decode_responses=True)
from collections import deque
from time import localtime, sleep, strftime, time, strptime
pressure_history = deque(maxlen=500)
_pubsub = redis_connection.pubsub()
_pubsub.subscribe(['bmp280', 'bme280', 'transfer_data'])
#ca 17/min with 0.05s
last_record = 0
last_log = 0
old_pressure = 0
last_transfer = None
log_altitude = False
log_pressure = False

for item in _pubsub.listen():
    if not item[u'type'] == 'message':
        continue
    if item['channel'] == 'transfer_data':
        _transfer_data = json.loads(item['data'])
        last_transfer = _transfer_data['utc']
        old_pressure = _transfer_data['pressure']
        continue
    data = json.loads(item['data'])
    pressure_history.append(data['pressure'])
    now = time()
    if last_record < now - 5:
        pressure_data = {
            'pressure': int(round(np.mean(pressure_history))),
            'temperature': round(data['temperature'], 1),
            'utc': int(round(now)),
            'hostname': data['hostname']}
        if item['channel'] == 'bme280':
            pressure_data['humidity'] = int(round(data['humidity']))
        if log_pressure and last_log < now - 60:
            key = 'pressure:{}:{}'.format(data['hostname'],
                                              strftime("%Y%m"))
            redis_connection.lpush(key, json.dumps(pressure_data))
            last_log = now
        redis_connection.set('current_pressure', json.dumps(pressure_data))
        last_record = now
    if log_altitude and np.abs(np.diff([data['pressure'], old_pressure])) > 5:
        key = 'altitude:{}:{}'.format(data['hostname'], strftime("%Y%m%d"))
        redis_connection.lpush(key, json.dumps(data))
        old_pressure = data['pressure']
    sleep(0.1)
    continue


import pip
pip.main(["install","matplotlib"])
import matplotlib.pyplot as plt
import numpy as np
import mxnet as mx
from mxnet import gluon, nd, image
from mxnet.gluon.data.vision import transforms
from gluoncv.data.transforms import video
from gluoncv import utils
from gluoncv.model_zoo import get_model
from gluoncv.utils import try_import_cv2
from ipywidgets import Video,Layout
import re

from gluoncv.utils.filesystem import try_import_decord

def TSN_PIC(img):
    img = image.imread(img)

    fig, ax = plt.subplots(figsize=(18, 18))
    ax.imshow(img.asnumpy())

    transform_fn = transforms.Compose([
    video.VideoCenterCrop(size=224),
    video.VideoToTensor(),
    video.VideoNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    img_list = transform_fn([img.asnumpy()])
    net = get_model('vgg16_ucf101', nclass=101, pretrained=True)
    pred = net(nd.array(img_list[0]).expand_dims(axis=0))
    classes = net.classes
    topK = 5
    ind = nd.topk(pred, k=topK)[0].astype('int')
    print('The input image is classified to be')
    results=[]
    for i in range(topK):
        datadict = {
             'header': ('\t%s'%(re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', ((classes[ind[i].asscalar()]).split('-',1)[0])))),
             'data': ('%.2f'%(nd.softmax(pred)[0][ind[i]].asscalar()*100)+'%')
        }
        results.append(datadict)
    print (results)
    return results
#        result = ('\t%s - %.2f'%(classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100)+'%')
#        result.replace(" ", "_")
#        result.replace(" ", "")
#        datadict = {
#             'data': result
#        }
#        results.append(datadict)
#    print (results)
#    return results
#    for i in range(topK):
#        results = []
#        result = ('\t%s - %.2f'%(classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#
#        datadict = {
#             'data': result
#        }
#        results.append(datadict.copy())
#    return results
   # return plt.show()
       #    print((result))
          #   'header':
              #print (results)

#TSN_PIC('static/playing_celo.jpeg')

global final_pred
def TSN_VIDEO(vid):
    cv2 = try_import_cv2()
    cap = cv2.VideoCapture(vid)
    cnt = 0
    video_frames = []
    transform_fn = transforms.Compose([
    video.VideoCenterCrop(size=224),
    video.VideoToTensor(),
    video.VideoNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    net = get_model('vgg16_ucf101', nclass=101, pretrained=True)
    while(cap.isOpened()):
        ret, frame = cap.read()
        cnt += 1
        if ret and cnt % 25 == 0:
            video_frames.append(frame)
        if not ret: break
    cap.release()
    if video_frames:
        video_frames_transformed = transform_fn(video_frames)
        final_pred = 0
        for _, frame_img in enumerate(video_frames_transformed):
            pred = net(nd.array(frame_img).expand_dims(axis=0))
            final_pred += pred
    final_pred /= len(video_frames)

    classes = net.classes
    topK = 5
    ind = nd.topk(final_pred, k=topK)[0].astype('int')
    print('The input video frame is classified to be')
    results=[]
    for i in range(topK):
        datadict = {
             'header': ('\t%s'%(re.sub(r'([a-z](?=[A-Z])|[A-Z](?=[A-Z][a-z]))', r'\1 ', (classes[ind[i].asscalar()]))).split('-',1)[0]),
             'data': ('%.2f'%(nd.softmax(pred)[0][ind[i]].asscalar()*100)+'%')
        }
        results.append(datadict)
    print (results)
    return results
#    results=[]
#    for i in range(topK):
#        results=('\t%s - %.2f'%(classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#        datadict = {
#          #   'header':
#             'data': results
#        }
#        results.append(datadict)
##        results=results+" "
#    #print (results)
#    return results
   # return Video.from_file(vid,layout=Layout(width='80%', height='80%'))

#TSN_VIDEO('static/v_Basketball_g01_c01.mp4')

def I3D_VID(vid):
    decord = try_import_decord()
    vr = decord.VideoReader(vid)
    frame_id_list = range(0, 64, 2)
    video_data = vr.get_batch(frame_id_list).asnumpy()
    clip_input = [video_data[vid, :, :, :] for vid, _ in enumerate(frame_id_list)]
    transform_fn = video.VideoGroupValTransform(size=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    clip_input = transform_fn(clip_input)
    clip_input = np.stack(clip_input, axis=0)
    clip_input = clip_input.reshape((-1,) + (32, 3, 224, 224))
    clip_input = np.transpose(clip_input, (0, 2, 1, 3, 4))
    model_name = 'i3d_inceptionv1_kinetics400'
    net = get_model(model_name, nclass=400, pretrained=True)
    pred = net(nd.array(clip_input))
    classes = net.classes
    topK = 5
    ind = nd.topk(pred, k=topK)[0].astype('int')
    print('The input video frame is classified to be')
    results=[]
    for i in range(topK):
#       result = ('\t%s - %.2f'%(((classes[ind[i].asscalar()]).replace("_", " ")).capitalize(), (nd.softmax(pred)[0][ind[i]].asscalar()*100))+'%')
#        result.replace("_", " ")
        datadict = {
             'header': ('\t%s'%(((classes[ind[i].asscalar()]).replace("_", " ")).title()).split('-',1)[0]),
             'data': ('%.2f'%(nd.softmax(pred)[0][ind[i]].asscalar()*100)+'%')
        }
        results.append(datadict)
    print (results)
    return results
#    results=[]
#    for i in range(topK):
#        results=('\t%s - %.2f'%(classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#        datadict = {
#          #   'header':
#             'data': results
#        }
#        results.append(datadict)
#    #print (results)
#    return results
#        result=('\t%s - %.2f'%
#          (classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#        print(result)
#    return Video.from_file(vid)

#I3D_VID('static/v_Basketball_g01_c01.mp4')

def SlowFast(vid):
    decord = try_import_decord()
    vr = decord.VideoReader(vid)
    fast_frame_id_list = range(0, 64, 2)
    slow_frame_id_list = range(0, 64, 16)
    frame_id_list = list(fast_frame_id_list) + list(slow_frame_id_list)
    video_data = vr.get_batch(frame_id_list).asnumpy()
    clip_input = [video_data[vid, :, :, :] for vid, _ in enumerate(frame_id_list)]
    transform_fn = video.VideoGroupValTransform(size=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    clip_input = transform_fn(clip_input)
    clip_input = np.stack(clip_input, axis=0)
    clip_input = clip_input.reshape((-1,) + (36, 3, 224, 224))
    clip_input = np.transpose(clip_input, (0, 2, 1, 3, 4))
    model_name = 'slowfast_4x16_resnet50_kinetics400'
    net = get_model(model_name, nclass=400, pretrained=True)
    pred = net(nd.array(clip_input))

    classes = net.classes
    topK = 5
    ind = nd.topk(pred, k=topK)[0].astype('int')
    print('The input video frame is classified to be')
    results=[]
    for i in range(topK):
        datadict = {
             'header': ('\t%s'%(((classes[ind[i].asscalar()]).replace("_", " ")).title()).split('-',1)[0]),
             'data': ('%.2f'%(nd.softmax(pred)[0][ind[i]].asscalar()*100)+'%')
        }
        results.append(datadict)
    print (results)
    return results
#    results=[]
#    for i in range(topK):
#        result=('\t%s - %.2f'%(classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#        datadict = {
#          #   'header':
#             'data': result
#        }
#        results.append(datadict)
#    #print (results)
#    return results
#        result=('\t%s - %.2f'%
#          (classes[ind[i].asscalar()], nd.softmax(pred)[0][ind[i]].asscalar()*100))
#        print(result)
#    return Video.from_file(vid,layout=Layout(width='80%', height='80%'))

#SlowFast('static/v_Basketball_g01_c01.mp4')


import Tkinter as tk
import numpy as np
import math
import time

class GUI(object):
    def __init__(self,environment_details,car_details,trace=False):
        self.w,self.h = 1280,720
        self.graphs = ['Average loss','Final score','Cross score']
        self.runs = ['Learn','Run only']
        self.construct_window()
        self.title_label('Reinforcement Learning Car Simulation')
        self.env = environment_details
        self.cars = [dict(car) for car in car_details]
        self.trace = trace
        self.trace_history_limit = 1500
        self.init_env()
        self.init_car()
        self.init_graph()
        self.approximator = lambda x: str(round(x, -int(math.floor(math.log10(abs(x))))+3)) if x!=0 else '0'
        self.refresh()

    def construct_window(self):
        # Level 1: Main window and frame
        self.window = tk.Tk()
        self.window.resizable(0,0)
        self.window_frame = tk.Frame(self.window,padx=10,pady=10)
        self.window_frame.pack()
        # Level 2: Segments for display and options
        proportion = 0.7
        self.display = tk.Canvas(self.window_frame, width=int(proportion*self.w), height=self.h)
        self.display.pack(side=tk.LEFT)
        self.display.config(background='gray90')
        self.options = tk.Frame(self.window_frame,padx=25,pady=25)
        self.options.pack(side=tk.RIGHT,fill=tk.Y)
        self.options.config(background='white smoke')
        # Level 3: Graph and debug info inside options. Display has no other children
        self.graph = tk.Canvas(self.options,background='white',width=int((1-proportion)*self.w), height=int(self.h/2))
        self.graph.pack(side=tk.BOTTOM)
        self.run_select = tk.StringVar()
        self.run_select_button_holder = tk.Frame(self.options)
        self.run_select_button_holder.pack(side=tk.TOP,anchor=tk.CENTER)
        self.debug_info = tk.StringVar()
        tk.Label(self.options,textvariable=self.debug_info,justify=tk.LEFT,background='white smoke',pady=10,font=('System')).pack(side=tk.TOP)
        self.graph_select = tk.StringVar()
        self.graph_select.set(self.graphs[0])
        self.graph_select_button_holder = tk.Frame(self.options)
        self.graph_select_button_holder.pack(side=tk.BOTTOM,anchor=tk.CENTER)
        # Level 4: Configure buttons
        for r in self.runs:
            tk.Radiobutton(self.run_select_button_holder, text=r, variable=self.run_select, value=r, indicatoron=0,font=('System')).pack(side=tk.LEFT,anchor=tk.CENTER)
        for g in self.graphs:
            tk.Radiobutton(self.graph_select_button_holder, text=g, variable=self.graph_select, value=g, indicatoron=0,font=('System')).pack(side=tk.LEFT,anchor=tk.CENTER)

    def set_display_range(self,xmin,xmax,ymin,ymax):
        x_range,y_range = xmax-xmin+2,ymax-ymin+2
        self.display_w, self.display_h = int(self.display.config()['width'][-1]),int(self.display.config()['height'][-1])
        self.scale_factor = min(self.display_w/x_range, self.display_h/y_range)
        self.center_offset = (self.scale_factor*2,self.scale_factor*2)

    def scale_and_offset_center(self,list_of_points):
        tr_pts = []
        for pt in list_of_points:
            tr_pts.append(int((pt[0]*self.scale_factor)+self.center_offset[0]))
            tr_pts.append(self.display_h-int((pt[1]*self.scale_factor)+self.center_offset[1]))
        return tr_pts

    def rotation_matrix(self,theta):
        ct = math.cos(theta)
        st = math.sin(theta)
        R = np.array([[ct,-st],[st,ct]])
        return R

    def init_env(self):
        self.set_display_range(min([a for (a,b) in self.env['points']]), max([a for (a,b) in self.env['points']]), min([b for (a,b) in self.env['points']]), max([b for (a,b) in self.env['points']]))
        track_coords = self.scale_and_offset_center(self.env['points'])
        self.track_id = self.display.create_polygon(track_coords,fill='white',outline='black')

    def init_car(self):
        for i in range(len(self.cars)):
            # Trace
            self.trace_mod(i,self.cars[i]['state'][0:2],init=True)
        for i in range(len(self.cars)):
            # Car
            R = self.rotation_matrix(self.cars[i]['state'][2])
            w = self.cars[i]['W']
            l = self.cars[i]['L']
            points = np.array([ [0,-w], [l,-w], [l,w], [0,w] ]).T
            points = np.dot(R,points)
            points[0,:] += self.cars[i]['state'][0]
            points[1,:] += self.cars[i]['state'][1]
            self.cars[i]['gui_body_id'] = self.display.create_polygon(self.scale_and_offset_center(points.T),fill='blue',outline='black')
            # Sensor
            for j in range(len(self.cars[i]['sensors'])):
                sa,sr = self.cars[i]['sensors'][j]['angle'],self.cars[i]['sensors'][j]['range']
                R = self.rotation_matrix(self.cars[i]['state'][2]+sa)
                points = np.array([ [0,0],[sr,0] ]).T
                points = np.dot(R,points)
                points[0,:] += self.cars[i]['state'][0]
                points[1,:] += self.cars[i]['state'][1]
                self.cars[i]['sensors'][j]['gui_sensor_id'] = self.display.create_line(self.scale_and_offset_center(points.T),fill='red')

    def init_graph(self):
        self.graph_w,self.graph_h = int(self.graph.config()['width'][-1]),int(self.graph.config()['height'][-1])
        self.plots = [self.graph.create_line([(0,0),(1,1)],width=2.0),
                        self.graph.create_text(0,self.graph_h,anchor=tk.SW,text=str((0,0))),
                        self.graph.create_text(0,0,anchor=tk.NW,text=str(1)),
                        self.graph.create_text(self.graph_w,self.graph_h,anchor=tk.SE,text=str(1))]

    def trace_mod(self,car_id,pt,init=False,force_end_line=False):
        if init==True:
            self.cars[car_id]['trace_history'] = [-1]*self.trace_history_limit
            self.cars[car_id]['trace_history_index'] = 0
            self.cars[car_id]['trace_history_buffer'] = []
        self.cars[car_id]['trace_history_buffer'].append(pt)
        if len(self.cars[car_id]['trace_history_buffer'])>100 or force_end_line==True:
            self.display.delete(self.cars[car_id]['trace_history'][self.cars[car_id]['trace_history_index']])
            if len(self.cars[car_id]['trace_history_buffer'])>2:
                self.cars[car_id]['trace_history'][self.cars[car_id]['trace_history_index']] = self.display.create_line(self.scale_and_offset_center(self.cars[car_id]['trace_history_buffer']),fill='gray65',width=1)
                self.cars[car_id]['trace_history_index'] += 1
                if self.cars[car_id]['trace_history_index'] >= self.trace_history_limit: self.cars[car_id]['trace_history_index'] = 0
            self.cars[car_id]['trace_history_buffer'] = []

    def update(self,car_id,state,draw_car=True,force_end_line=False):
        if self.trace==True:
            self.trace_mod(car_id,state[0:2],force_end_line=force_end_line)
        if draw_car==True:
            R = self.rotation_matrix(state[2])
            w = self.cars[car_id]['W']
            l = self.cars[car_id]['L']
            points = np.array([ [0,-w], [l,-w], [l,w], [0,w] ]).T
            points = np.dot(R,points)
            points[0,:] += state[0]
            points[1,:] += state[1]
            translated_points = self.scale_and_offset_center(points.T)
            self.display.coords(self.cars[car_id]['gui_body_id'], *translated_points)
            self.display.tag_raise(self.cars[car_id]['gui_body_id'])
            for j in range(len(self.cars[car_id]['sensors'])):
                sa,sr = self.cars[car_id]['sensors'][j]['angle'],self.cars[car_id]['sensors'][j]['range']
                R = self.rotation_matrix(state[2]+sa)
                points = np.array([ [0,0],[sr,0] ]).T
                points = np.dot(R,points)
                points[0,:] += state[0]
                points[1,:] += state[1]
                translated_points = self.scale_and_offset_center(points.T)
                self.display.coords(self.cars[car_id]['sensors'][j]['gui_sensor_id'],*translated_points)
                self.display.tag_raise(self.cars[car_id]['sensors'][j]['gui_sensor_id'])

    def update_graph(self,X,Y,plt_idx=None,color='black',thickness=1.0,style=None):
        if plt_idx != self.graph_select.get(): return
        if len(X)<2: return
        minX,minY,maxX,maxY = min(X),min(Y),max(X),max(Y)
        w,h = maxX-minX, maxY-minY
        yscale = minY if h==0 else float(self.graph_h)/h
        xscale = minX if w==0 else float(self.graph_w)/w
        scaled_pts = []
        for (x,y) in zip(X,Y):
            scaled_pts.append(int(xscale*(x-minX)))
            scaled_pts.append(self.graph_h - int(yscale*(y-minY)))
        self.graph.coords(self.plots[0],*scaled_pts)
        self.graph.itemconfig(self.plots[1],text=self.approximator(minX)+','+self.approximator(minY))
        self.graph.itemconfig(self.plots[2],text=self.approximator(maxY))
        self.graph.itemconfig(self.plots[3],text=self.approximator(maxX))

    def update_debug_info(self,info):
        self.debug_info.set(info)

    def refresh(self):
        self.window.update()
        self.window.update_idletasks()
        time.sleep(0.005)

    def enable_trace(self):
        self.trace = True

    def disable_trace(self):
        self.trace = False

    def title_label(self,label):
        self.window.title(label)

    def get_run_select(self):
        return self.run_select.get()

    def set_run_select(self,value):
        self.run_select.set(value)

    def remove_traces(self):
        for car in self.cars:
            car['trace_history_index'] = 0
            car['trace_history_buffer'] = []
            for idx in range(len(car['trace_history'])):
                self.display.delete(car['trace_history'][idx])


import json
import numpy
from functools import reduce

class RicoGesture:
    TYPE_PRESS = 0
    TYPE_DRAG = 1

    def __init__(self, type, activity, start, end = [0, 0]):
        self.type = type
        self.activity = activity

        self.start_x = int(start[0])
        self.start_y = int(start[1])

        self.end_x = int(end[0])
        self.end_y = int(end[1])

    def __str__(self):
        return "<RicoGesture\tType: %d, Start: [%d,%d], End: [%d, %d]>" % (self.type, self.start_x, self.start_y, self.end_x, self.end_y)

class RicoGestureUtils:

    def __init__(self, size_x, size_y, base_path):
        self.size = [size_x, size_y]
        self.gestures = {}
        self.base_path = base_path

    def reduce_gestures_updating_size(self, acc, elem):
        acc[elem] = {"gestures": numpy.multiply(self.size, self.gestures[elem])}
        return acc

    def reduce_gestures_updating_activity(self, acc, elem):
        with open("%s/view_hierarchies/%s.json" % (self.base_path, elem), "r") as file:
            hier = json.load(file)
            try:
                acc[elem]["activity"] = (hier["activity_name"].split("/"))[1]
            except Exception as e:
                acc[elem]["activity"] = "null"
        #acc[elem]["activity"]
        return acc

    def parseGesturesFile(self):
        with open("%s/gestures.json" % self.base_path, "r") as g_file:
            self.gestures = json.load(g_file)
            self.gestures = {int(k): v for k,v in self.gestures.items()}
        self.gestures = reduce(lambda acc, elem: self.reduce_gestures_updating_size(acc, elem), self.gestures.keys(), {})
        self.gestures = reduce(lambda acc, elem: self.reduce_gestures_updating_activity(acc, elem), self.gestures.keys(), self.gestures)

    def getNextGesture(self):
        for gesture_key in sorted(self.gestures.keys()):
            gesture = self.gestures[gesture_key]["gestures"]
            #gesture = numpy.multiply(self.size, gesture)
            if len(gesture) > 0:
                if len(gesture) == 1:
                    yield RicoGesture(RicoGesture.TYPE_PRESS, self.gestures[gesture_key]["activity"], gesture[0])
                else:
                    yield RicoGesture(RicoGesture.TYPE_DRAG, self.gestures[gesture_key]["activity"], gesture[0], gesture[-1])
        raise StopIteration()


#!/usr/bin/python
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
import sys
import os
import re
import io
import commands
import logging
import time
import numpy as np
import thread
import paddle
import paddle.fluid as fluid
import paddle.fluid.incubate.fleet.base.role_maker as role_maker
from paddle.fluid.incubate.fleet.parameter_server.distribute_transpiler import fleet
from paddle.fluid.transpiler.distribute_transpiler import DistributeTranspilerConfig
import reader_generator as py_reader
from paddle.fluid.contrib.utils import HDFSClient

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger("fluid")
logger.setLevel(logging.INFO)

class FleetDistRunnerBase(object):
    """
    Distribute training base class:
        This class abstracts the training process into several major steps:
        1. input_data: input data of network, this function should be realized by user
        2. net: network definition, this function should be defined by user
        3. run_pserver: run pserver node in distribute environment
        4. run_trainer: run trainer, choose the way of training network according to requirement params
        5. run_infer: prediction based on the trained model
        6. py_reader: using py_reader method get data, this function should be realized by user
        7. dataset_reader: using dataset method get data, this function should be realized by user
        8. runtime_main: program entry, get the environment parameters, decide which function to call
    """

    def input_data(self, params):
        """
        Function input_data: Definition of input data format in the network
        Args:
            :params: the hyper parameters of network
        Returns:
            defined by users
        """
        raise NotImplementedError(
            "input_data should be implemented by child classes.")

    def net(self, inputs, params):
        """
        Function net: Definition of network structure
        Args:
            :inputs: input data, eg: dataset and labels. defined by funtion: self.input_data
            :params: the hyper parameters of network
        Returns:
            evaluation parameter, defined by users
        """
        raise NotImplementedError(
            "net should be implemented by child classes.")

    def infer_net(self, params):
        """
        Function net: Definition of infer network structure, This function is not required
                      if the prediction is same with the training logic
        Args:
            :params: the hyper parameters of network
        Returns:
            evaluation parameter, defined by users
         """
        raise NotImplementedError(
            "net should be implemented by child classes.")

    def print_pserver_learning_rate(self, params):
        while(True):
            step = fluid.global_scope().find_var("@LR_DECAY_COUNTER@").get_tensor()
            lr_ = fluid.global_scope().find_var("tmp_3").get_tensor()
            logger.info(np.array(step))
            logger.info(np.array(lr_))
            time.sleep(10)

    def run_pserver(self, params):
        """
        Function run_pserver: Operation method of parameter server
        Args
            :params the hyper parameters of network
        Returns:
            None
        """
        # step1: define the role of node, configure communication parameter
        role = role_maker.UserDefinedRoleMaker(
            current_id=params.current_id,
            role=role_maker.Role.SERVER,
            worker_num=params.trainers,
            server_endpoints=params.pserver_endpoints)
        fleet.init(role)

        # step2: define the input data of network
        # define the model, build the pserver program
        inputs = self.input_data(params)
        if params.is_pyreader_train:
            reader = self.py_reader(params)
            inputs = fluid.layers.read_file(reader)
        elif not params.is_dataset_train:
            raise ValueError("Program must has Date feed method: is_pyreader_train / is_dataset_train")

        # step3: define the network
        loss = self.net(inputs, params)

        # step4: define the optimizer for your model
        optimizer = fluid.optimizer.SGD(
            learning_rate=fluid.layers.exponential_decay(
                learning_rate=params.learning_rate,
                decay_steps=params.decay_steps,
                decay_rate=params.decay_rate,
                staircase=True))
        optimizer = fleet.distributed_optimizer(optimizer, self.strategy)
        optimizer.minimize(loss)

        fleet.init_server()
        logger.info("PServer init success!")

        with open("pserver_train.proto", 'w') as f:
            f.write(str(fleet.main_program))
        with open("pserver_startup.proto", 'w') as f:
            f.write(str(fleet.startup_program))
        #try:
        #    thread.start_new_thread(self.print_pserver_learning_rate, (params, ))
        #except:
        #    logger.info("Error: unable to start thread")
        fleet.run_server()

    def run_dataset_trainer(self, params):
        """
        Function run_dataset_trainer: Operation method of training node
        Args:
            :params params: the hyper parameters of network
        Returns
            :train_result: the dict of training log
        """
        # step1: define the role of node, configure communication parameter
        role = role_maker.UserDefinedRoleMaker(
            current_id=params.current_id,
            role=role_maker.Role.WORKER,
            worker_num=params.trainers,
            server_endpoints=params.pserver_endpoints)
        fleet.init(role)

        # step2: define the input data of network
        inputs = self.input_data(params)

        # step3: define the network, same with PSERVER
        # For the model: word2vec, we use loss to measure the performance of network
        # Replace it with your network evaluation index
        loss = self.net(inputs, params)

        # step4: define the optimizer for your model
        optimizer = fluid.optimizer.SGD(
            learning_rate=fluid.layers.exponential_decay(
                learning_rate=params.learning_rate,
                decay_steps=params.decay_steps,
                decay_rate=params.decay_rate,
                staircase=True))
        optimizer = fleet.distributed_optimizer(optimizer, self.strategy)
        optimizer.minimize(loss)

        # step5: define Executor and run startup program
        exe = fluid.Executor(fluid.CPUPlace())
        fleet.init_worker()
        exe.run(fleet.startup_program)

        with open(str(params.current_id) + "_trainer_train.proto", 'w') as f:
            f.write(str(fleet.main_program))

        with open(str(params.current_id) + "_trainer_startup.proto", 'w') as f:
            f.write(str(fleet.startup_program))

        # step6: init dataset reader
        # Notice: Both dataset and py_reader method don't using feed={dict} to input data
        # Paddle Fluid enter data by variable name
        # When we do the definition of the reader, the program has established the workflow
        dataset = self.dataset_reader(inputs, params)
        file_list = [str(params.train_files_path) + "/%s" % x
                     for x in os.listdir(params.train_files_path)]
        if params.is_local_cluster:
            file_list = fleet.split_files(file_list)
        logger.info("file list: {}".format(file_list))
        logger.info("there are a total of {} files.".format(len(file_list)))
        logger.info('----------------------NO.%s trainer ready----------------' % (params.current_id))
        all_examples = self.get_example_num(file_list)

        # step7: begin to train your model, good luck
        train_result = {}
        for epoch in range(params.epochs):
            dataset.set_filelist(file_list)
            start_time = time.time()
            # Notice: function train_from_dataset does not return fetch value
            exe.train_from_dataset(program=fleet.main_program, dataset=dataset,
                                   fetch_list=[loss], fetch_info=['loss'],
                                   print_period=1000, debug=False)
            end_time = time.time()
            speed = float(all_examples) / float(end_time - start_time)
            logger.info("epoch: %d finished, speed: %f" % (epoch, speed))

            self.record_speed(epoch, train_result, speed)
            self.record_memory(epoch, train_result)
            if params.is_first_trainer and params.test:
                model_path = str(params.model_path) + '/trainer_' + str(params.current_id) + '_epoch_' + str(epoch)
                fleet.save_persistables(executor=exe, dirname=model_path)
                if not params.is_local_cluster:
                    self.upload_files(model_path, params, 'model')

        log_path = str(params.log_path + '/' + str(params.current_id) + '_dataset_train.log')
        with open(log_path, 'w+') as f:
            f.write(str(train_result))
        if not params.is_local_cluster:
            self.upload_files(log_path, params, 'log')

        logger.info("Train Success!")
        fleet.stop_worker()
        return train_result

    def run_pyreader_trainer(self, params):
        """
        Function run_trainer: Operation method of training node
        Args:
            :params params: the hyper parameters of network
        Returns
            :train_result: the dict of training log
        """
        # step1: define the role of node, configure communication parameter
        role = role_maker.UserDefinedRoleMaker(
            current_id=params.current_id,
            role=role_maker.Role.WORKER,
            worker_num=params.trainers,
            server_endpoints=params.pserver_endpoints)
        fleet.init(role)

        file_list = [str(params.train_files_path) + "/%s" % x
                     for x in os.listdir(params.train_files_path)]
        if params.is_local_cluster:
            file_list = fleet.split_files(file_list)
        logger.info("file list: {}".format(file_list))
        logger.info("there are a total of {} files.".format(len(file_list)))
        word2vec_reader = py_reader.Word2VecReader(params.dict_path, file_list, 0, 1)
        params.dict_size = word2vec_reader.dict_size

        # step2: define the input data of network
        inputs = self.input_data(params)
        reader = self.py_reader(params)
        inputs = fluid.layers.read_file(reader)

        # step3: define the network, same with PSERVER
        # For the model: word2vec, we use loss to measure the performance of network
        # Replace it with your network evaluation index,
        loss = self.net(inputs, params)

        # step4: define the optimizer for your model
        # define the optimizer for your model
        lr = fluid.layers.exponential_decay(
                learning_rate=params.learning_rate,
                decay_steps=params.decay_steps,
                decay_rate=params.decay_rate,
                staircase=True)

        optimizer = fluid.optimizer.SGD(learning_rate=lr)
        optimizer = fleet.distributed_optimizer(optimizer, self.strategy)
        optimizer.minimize(loss)

        # step5: define Executor and run startup program
        exe = fluid.Executor(fluid.CPUPlace())
        fleet.init_worker()
        exe.run(fleet.startup_program)
        CPU_NUM = int(params.cpu_num)
        with open(str(params.current_id) + "_trainer_train.proto", 'w') as f:
            f.write(str(fleet.main_program))

        with open(str(params.current_id) + "_trainer_startup.proto", 'w') as f:
            f.write(str(fleet.startup_program))

        # step6: init py_reader reader
        # Notice: Both dataset and py_reader method don't using feed={dict} to input data
        # Paddle Fluid enter data by variable name
        # When we do the definition of the reader, the program has established the workflow

        np_power = np.power(np.array(word2vec_reader.id_frequencys), 0.75)
        id_frequencys_pow = np_power / np_power.sum()
        reader.decorate_tensor_provider(
            py_reader.convert_python_to_tensor(id_frequencys_pow, params.batch_size, word2vec_reader.train(), params))

        # step7: define the compiled program
        exec_strategy = fluid.ExecutionStrategy()
        exec_strategy.num_threads = int(params.cpu_num)
        exec_strategy.use_experimental_executor = True

        build_strategy = fluid.BuildStrategy()
        build_strategy.async_mode = self.async_mode
        if CPU_NUM > 1:
            build_strategy.reduce_strategy = fluid.BuildStrategy.ReduceStrategy.Reduce

        compiled_prog = fluid.compiler.CompiledProgram(
            fleet.main_program).with_data_parallel(
            loss_name=loss.name, build_strategy=build_strategy, exec_strategy=exec_strategy)
        logger.info('----------------------NO.%s trainer ready----------------' % (params.current_id))

        # step8: begin to train your model, good luck
        train_result = {}
        all_examples = self.get_example_num(file_list)
        for epoch in range(params.epochs):
            reader.start()
            start_time = time.time()
            batch_id = 0
            # py_reader need use "try & catch Exception" method to load data continuously
            try:
                while True:
                    loss_val = exe.run(program=compiled_prog, fetch_list=[loss.name])
                    loss_val = np.mean(loss_val)

                    if batch_id % 1000 == 0 and batch_id != 0:
                        logger.info(
                            "TRAIN --> pass: {} batch: {} loss: {} reader queue:{}".
                                format(epoch, batch_id, loss_val.mean(), reader.queue.size()))
                    batch_id += 1
            except fluid.core.EOFException:
                reader.reset()
            end_time = time.time()
            speed = float(all_examples) / float(end_time - start_time)
            logger.info("epoch: %d finished, speed: %f" % (epoch, speed))

            train_result = self.record_speed(epoch, train_result, speed)
            train_result = self.record_memory(epoch, train_result)

            if params.is_first_trainer and params.test:
                model_path = str(params.model_path) + '/trainer_' + str(params.current_id) + '_epoch_' + str(epoch)
                fleet.save_persistables(executor=exe, dirname=model_path)
                if not params.is_local_cluster:
                    self.upload_files(model_path, params, 'model')

        log_path = str(params.log_path + '/' + str(params.current_id) + '_pyreader_train.log')
        with open(log_path, 'w+') as f:
            f.write(str(train_result))
        if not params.is_local_cluster:
            self.upload_files(log_path, params, 'log')

        logger.info("Train Success!")
        fleet.stop_worker()
        return train_result

    def get_example_num(self,file_list):
        count = 0
        for f in file_list:
            last_count = count
            for index, line in enumerate(open(f, 'r')):
                count += 1
            logger.info("file: %s has %s examples"%(f,count-last_count))
        logger.info("Total example: %s"%count)
        return count

    def upload_files(self, local_path, params, kind):
        """
        upload files to hdfs
        """
        import paddlecloud.upload_utils as upload_utils
        sys_job_id = os.getenv("SYS_JOB_ID")
        output_path = os.getenv("OUTPUT_PATH")
        remote_path = output_path + "/" + sys_job_id + "/" + kind + "/"
        if (not params.is_local_cluster) and params.test:
            upload_rst = upload_utils.upload_to_hdfs(local_file_path=local_path, remote_file_path=remote_path)
            logger.info("remote_path: {}, upload_rst: {}".format(remote_path, upload_rst))

    def run_local_reader(self, params):
        place = fluid.CPUPlace()
        inputs = self.input_data(params)
        reader = self.py_reader(params)
        inputs = fluid.layers.read_file(reader)

        file_list = [str(params.train_files_path) + "/%s" % x
                     for x in os.listdir(params.train_files_path)]
        logger.info("file list: {}".format(file_list))
        logger.info("there are a total of {} files.".format(len(file_list)))
        word2vec_reader = py_reader.Word2VecReader(params.dict_path, file_list, 0, 1)
        params.dict_size = word2vec_reader.dict_size

        np_power = np.power(np.array(word2vec_reader.id_frequencys), 0.75)
        id_frequencys_pow = np_power / np_power.sum()
        reader.decorate_tensor_provider(py_reader.convert_python_to_tensor(id_frequencys_pow, params.batch_size, word2vec_reader.train(), params))

        loss = self.net(inputs, params)

        optimizer = fluid.optimizer.SGD(
            learning_rate=fluid.layers.exponential_decay(
                learning_rate=params.learning_rate,
                decay_steps=params.decay_steps,
                decay_rate=params.decay_rate,
                staircase=True))
        optimizer.minimize(loss)
        train_program = fluid.default_main_program()

        train_result = {}
        exe = fluid.Executor(place)
        exe.run(fluid.default_startup_program())

        exec_strategy = fluid.ExecutionStrategy()
        exec_strategy.use_experimental_executor = True

        print("CPU_NUM:" + str(os.getenv("CPU_NUM", '1')))
        exec_strategy.num_threads = int(os.getenv("CPU_NUM", '1'))

        build_strategy = fluid.BuildStrategy()
        if int(os.getenv("CPU_NUM", '1')) > 1:
            build_strategy.reduce_strategy = fluid.BuildStrategy.ReduceStrategy.Reduce

        train_exe = fluid.ParallelExecutor(
            use_cuda=False,
            loss_name=loss.name,
            main_program=train_program,
            build_strategy=build_strategy,
            exec_strategy=exec_strategy)

        train_result = {}
        all_examples = self.get_example_num(file_list)
        logger.info("--------begin------- ")
        for epoch in range(params.epochs):
            reader.start()
            start_time = time.time()
            epoch_loss = 0.0
            total_time = 0.0
            batch_id = 0
            try:
                while True:
                    loss_val = train_exe.run(fetch_list=[loss.name])
                    loss_val = np.mean(loss_val)
                    epoch_loss += loss_val
                    if batch_id % 1000 == 0 and batch_id != 0:
                        logger.info(
                            "TRAIN --> pass: {} batch: {} loss: {} reader queue:{}".
                                format(epoch, batch_id, loss_val.mean(), reader.queue.size()))
                    batch_id += 1
            except fluid.core.EOFException:
                reader.reset()
            end_time = time.time()
            speed = float(all_examples) / float(end_time - start_time)
            logger.info("epoch: %d finished, speed: %f" % (epoch, speed))
            train_result = self.record_speed(epoch, train_result, speed)
            train_result = self.record_memory(epoch, train_result)
            train_result[epoch]['loss'] = epoch_loss / float(batch_id)
            sys.stderr.write("epoch %d finished, speed=%f examples/s, loss=%f\n"
                       % ((epoch + 1), speed, train_result[epoch]['loss']))
            model_path = str(params.model_path) +'/trainer_0_' + 'epoch_'+ str(epoch)
            fluid.io.save_persistables(executor=exe, dirname=model_path)

            log_path = str(params.log_path + '/' + str(epoch) + '.log')
            with open(log_path, 'w+') as f:
                f.write(str(train_result))
        print("Train Success!")

    def check_model_format(self, epoch_id):
        pattern = '^trainer_[0-9]+_epoch_[0-9]+$'
        if re.match(pattern, epoch_id, re.M|re.I):
            return True
        else:
            return False

    def run_infer(self, params, model_path):
        """
        Function run_infer: Operation method of prediction
        Args:
            :params params: the hyper parameters of network
        Returns
            :infer_result, type:dict, record the evalution parameter and program resource usage situation
        """
        if not os.path.isdir(model_path):
            logger.info("{} is not a dir".format(model_path))
            return
        epoch_id = os.path.basename(model_path)
        if not self.check_model_format(epoch_id):
            return
        if os.path.exists(model_path + '.acc'):
            return

        params.dict_size, test_reader, id2word = py_reader.prepare_data(
            params.test_files_path, params.infer_dict_path, batch_size=params.infer_batch_size)

        place = fluid.CPUPlace()
        exe = fluid.Executor(place)
        infer_result = {}
        startup_program = fluid.framework.Program()
        test_program = fluid.framework.Program()

        with fluid.framework.program_guard(test_program,startup_program):
            values, pred = self.infer_net(params)
            logger.info(model_path)
            fluid.io.load_persistables(
                executor=exe, dirname=model_path, main_program=fluid.default_main_program())

            accum_num = 0
            accum_num_sum = 0.0
            step_id = 0
            for data in test_reader():
                step_id += 1
                b_size = len([dat[0] for dat in data])
                wa = np.array([dat[0] for dat in data]).astype("int64").reshape(b_size, 1)
                wb = np.array([dat[1] for dat in data]).astype("int64").reshape(b_size, 1)
                wc = np.array([dat[2] for dat in data]).astype("int64").reshape(b_size, 1)

                label = [dat[3] for dat in data]
                input_word = [dat[4] for dat in data]
                para = exe.run(fluid.default_main_program(),
                               feed={
                                   "analogy_a": wa, "analogy_b": wb, "analogy_c": wc,
                                   "all_label":
                                       np.arange(params.dict_size).reshape(
                                           params.dict_size, 1).astype("int64"),
                               },
                               fetch_list=[pred.name, values],
                               return_numpy=False)
                pre = np.array(para[0])
                val = np.array(para[1])
                for ii in range(len(label)):
                    top4 = pre[ii]
                    accum_num_sum += 1
                    for idx in top4:
                        if int(idx) in input_word[ii]:
                            continue
                        if int(idx) == int(label[ii][0]):
                            accum_num += 1
                        break
                if step_id % 1 == 0:
                    logger.info("step:%d %d " % (step_id, accum_num))
            acc = 1.0 * accum_num / accum_num_sum
            logger.info("acc:%.3f " % acc)
            infer_result['acc'] = acc
        with open(model_path + '.acc', 'w') as fout:
            fout.write(str(infer_result) + '\n')
        #return infer_result

    def py_reader(self, params):
        """
        Function py_reader: define the data read method by fluid.layers.py_reader
        help: https://www.paddlepaddle.org.cn/documentation/docs/zh/1.5/api_cn/layers_cn/io_cn.html#py-reader
        Args:
            :params params: the hyper parameters of network
        Returns:
            defined by user
        """
        raise NotImplementedError(
            "py_reader should be implemented by child classes.")

    def dataset_reader(self, inputs, params):
        """
        Function dataset_reader: define the data read method by fluid.dataset.DatasetFactory
        help: https://www.paddlepaddle.org.cn/documentation/docs/zh/1.5/api_cn/dataset_cn.html#fluid-dataset
        Args:
           :params inputs: input data, eg: dataset and labels. defined by funtion: self.input_data
           :params params: the hyper parameters of network
        Returns:
           defined by user
        """
        raise NotImplementedError(
            "dataset_reader should be implemented by child classes.")

    def record_speed(self, epoch, train_result, speed):
        """
        record the operation speed
        """
        train_result[epoch] = {}
        train_result[epoch]['speed'] = speed
        return train_result

    def record_memory(self, epoch, train_result):
        info = process_info()
        logger.info(info)
        train_result[epoch]['memory'] = info['mem']
        train_result[epoch]['cpu'] = info['cpu']
        train_result[epoch]['rss'] = info['rss']
        train_result[epoch]['vsa'] = info['vsa']
        return train_result

    def check_all_trainers_ready(self, epoch):
        output = os.getenv("OUTPUT_PATH")
        trainer_num = int(os.getenv("PADDLE_TRAINERS_NUM", "1"))
        trainer_id = int(os.getenv("PADDLE_TRAINER_ID", "0"))
        user_id = os.getenv("SYS_USER_ID")
        job_id = os.getenv("SYS_JOB_ID")

        hadoop_home = os.getenv("HADOOP_HOME")
        configs = {
            "fs.default.name": os.getenv("FS_NAME"),
            "hadoop.job.ugi": os.getenv("FS_UGI")
        }

        node_ready = "ready.{}.{}.done".format(epoch, trainer_id)

        with open(node_ready, "w") as node:
            node.write("")

        ready_path = "{}/{}/{}/ready".format(output, user_id, job_id)

        client = HDFSClient(hadoop_home, configs)
        client.makedirs(ready_path)
        client.upload(hdfs_path=ready_path, local_path=node_ready, overwrite=True, retry_times=0)

        print("PUT {} ON HDFS {} OK".format(node_ready, ready_path))

        ready_num = len(client.ls(ready_path))
        while ready_num % trainer_num != 0:
            print("have {} trainers need to be ready".format(trainer_num - ready_num % trainer_num))
            time.sleep(10)
            ready_num = len(client.ls(ready_path))

        print("All trainers are ready, continue training")

    def runtime_main(self, params):
        """
        Function runtime_main: the entry point for program running
        Args:
            :params params: the hyper parameters of network
        """
        if params.is_local:
            logger.info("local train start")
            self.run_local_reader(params)
        else:
            logger.info("distributed train start")
            id_counts = []
            with io.open(params.dict_path, 'r', encoding='utf-8') as f:
                for line in f:
                    word, count = line.split()[0], int(line.split()[1])
                    id_counts.append(count)
            params.dict_size = len(id_counts)
            # Step1: get the environment variable, mainly related to network communication parameters
            params.role = os.getenv("TRAINING_ROLE")
            logger.info("Training role: {}".format(params.role))

            params.current_id = int(os.getenv("PADDLE_TRAINER_ID"))
            logger.info("Current Id: {}".format(params.current_id))

            params.trainers = int(os.getenv("PADDLE_TRAINERS_NUM"))
            logger.info("Trainer num: {}".format(params.trainers))

            params.pserver_ports = os.getenv("PADDLE_PORT")
            logger.info("Pserver ports: {}".format(params.pserver_ports))

            params.pserver_ip = os.getenv("PADDLE_PSERVERS")
            logger.info("Pserver IP: {}".format(params.pserver_ip))

            params.current_endpoint = os.getenv("POD_IP", "localhost") + ":" + params.pserver_ports

            params.cpu_num = os.getenv("CPU_NUM", "1")
            logger.info("output path: {}".format(params.model_path))

            # Step2: decide communication mode between PSERVER & TRAINER
            # recommended mode: pyreader + sync_mode / dataset + async_mode
            self.strategy = DistributeTranspilerConfig()
            if params.sync_mode == 'sync':
                self.strategy.sync_mode = True
                self.strategy.runtime_split_send_recv = False
                self.async_mode = False
                params.batch_size = int(params.batch_size / (params.trainers * int(params.cpu_num)))
            elif params.sync_mode == 'async':
                self.strategy.sync_mode = False
                self.async_mode = True
                self.strategy.runtime_split_send_recv = True
            elif params.sync_mode == "geo_async":
                self.strategy.sync_mode = False
                self.async_mode = True
                self.strategy.runtime_split_send_recv = True
                self.strategy.geo_sgd_mode = True
                self.strategy.geo_sgd_need_push_nums = 400
                params.decay_steps = int(int(params.decay_steps) / params.trainers)

            # Step3: Configure communication IP and ports
            if params.is_local_cluster:
                for port in params.pserver_ports.split(","):
                    logger.info("current port: %s" % port)
                    params.pserver_endpoints.append(':'.join(
                        [params.pserver_ip, port]))
                    logger.info("add pserver_endpoint:%s" % (params.pserver_endpoints))
            else:
                for ip in params.pserver_ip.split(","):
                    params.pserver_endpoints.append(':'.join(
                        [ip, params.pserver_ports]))

            params.endpoints = ",".join(params.pserver_endpoints)
            logger.info("pserver_endpoints: {}".format(params.pserver_endpoints))

            if params.role == "TRAINER" and params.current_id == 0:
                params.is_first_trainer = True

            # Step4: According to the parameters-> TRAINING_ROLE, decide which method to run
            train_result = {}
            if params.role == "PSERVER":
                self.run_pserver(params)
            elif params.role == "TRAINER":
                self.check_all_trainers_ready(0)
                if params.is_dataset_train:
                    train_result = self.run_dataset_trainer(params)
                elif params.is_pyreader_train:
                    train_result = self.run_pyreader_trainer(params)
            else:
                raise ValueError("Please choice training role for current node : PSERVER / TRAINER")

            logger.info("Distribute train success!")

def process_info():
    pid = os.getpid()
    res = commands.getstatusoutput('ps aux|grep ' + str(pid))[1].split('\n')[0]

    p = re.compile(r'\s+')
    l = p.split(res)
    info = {'user': l[0],
            'pid': l[1],
            'cpu': l[2],
            'mem': l[3],
            'vsa': l[4],
            'rss': l[5], }
    return info