from typing import Union
import torch
import numpy as np
from dqc.utils.datastruct import ZType

periodic_table_atomz = {
    "H": 1,
    "He": 2,
    "Li": 3,
    "Be": 4,
    "B": 5,
    "C": 6,
    "N": 7,
    "O": 8,
    "F": 9,
    "Ne": 10,
    "Na": 11,
    "Mg": 12,
    "Al": 13,
    "Si": 14,
    "P": 15,
    "S": 16,
    "Cl": 17,
    "Ar": 18,
    "K": 19,
    "Ca": 20,
    "Sc": 21,
    "Ti": 22,
    "V": 23,
    "Cr": 24,
    "Mn": 25,
    "Fe": 26,
    "Co": 27,
    "Ni": 28,
    "Cu": 29,
    "Zn": 30,
    "Ga": 31,
    "Ge": 32,
    "As": 33,
    "Se": 34,
    "Br": 35,
    "Kr": 36,
    "Rb": 37,
    "Sr": 38,
    "Y": 39,
    "Zr": 40,
    "Nb": 41,
    "Mo": 42,
    "Tc": 43,
    "Ru": 44,
    "Rh": 45,
    "Pd": 46,
    "Ag": 47,
    "Cd": 48,
    "In": 49,
    "Sn": 50,
    "Sb": 51,
    "Te": 52,
    "I": 53,
    "Xe": 54,
}

atom_masses = {  # isotope-averaged atom masses in a.m.u.
    # from https://www.angelo.edu/faculty/kboudrea/periodic/structure_mass.htm
    1: 1.00797,
    2: 4.00260,
    3: 6.941,
    4: 9.01218,
    5: 10.81,
    6: 12.011,
    7: 14.0067,
    8: 15.9994,
    9: 18.998403,
    10: 20.179,
    11: 22.98977,
    12: 24.305,
    13: 26.98154,
    14: 28.0855,
    15: 30.97376,
    16: 32.06,
    17: 35.453,
    18: 39.948,
    19: 39.0983,
    20: 40.08,
    21: 44.9559,
    22: 47.90,
    23: 50.9415,
    24: 51.996,
    25: 54.9380,
    26: 55.847,
    27: 58.9332,
    28: 58.70,
    29: 63.546,
    30: 65.38,
    31: 69.72,
    32: 72.59,
    33: 74.9216,
    34: 78.96,
    35: 79.904,
    36: 83.80,
    37: 85.4678,
    38: 87.62,
    39: 88.9059,
    40: 91.22,
    41: 92.9064,
    42: 95.94,
    43: 98.,
    44: 101.07,
    45: 102.9055,
    46: 106.4,
    47: 107.868,
    48: 112.41,
    49: 114.82,
    50: 118.69,
    51: 121.75,
    53: 126.9045,
    52: 127.60,
    54: 131.30,
}

# JCP 41, 3199 (1964); DOI:10.1063/1.1725697
# taken from PySCF:
# https://github.com/pyscf/pyscf/blob/45582e915e91890722fcae2bc30fb04867d5c95f/pyscf/data/radii.py#L23
# I don't know why H has 0.35 while in the reference it is 0.
# They are in angstrom, so we need to convert it to Bohr
atom_bragg_radii = list(np.array([
    2.00,                                                        # Ghost atom
    0.35,                                     1.40,              # 1s
    1.45, 1.05, 0.85, 0.70, 0.65, 0.60, 0.50, 1.50,              # 2s2p
    1.80, 1.50, 1.25, 1.10, 1.00, 1.00, 1.00, 1.80,              # 3s3p
    2.20, 1.80,                                                  # 4s
    1.60, 1.40, 1.35, 1.40, 1.40, 1.40, 1.35, 1.35, 1.35, 1.35,  # 3d
                1.30, 1.25, 1.15, 1.15, 1.15, 1.90,              # 4p
    2.35, 2.00,                                                  # 5s
    1.80, 1.55, 1.45, 1.45, 1.35, 1.30, 1.35, 1.40, 1.60, 1.55,  # 4d
                1.55, 1.45, 1.45, 1.40, 1.40, 2.10,              # 5p
    2.60, 2.15,                                                  # 6s
    1.95, 1.85, 1.85, 1.85, 1.85, 1.85, 1.85,                    # La, Ce-Eu
    1.80, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75,              # Gd, Tb-Lu
          1.55, 1.45, 1.35, 1.35, 1.30, 1.35, 1.35, 1.35, 1.50,  # 5d
                1.90, 1.80, 1.60, 1.90, 1.45, 2.10,              # 6p
    1.80, 2.15,                                                  # 7s
    1.95, 1.80, 1.80, 1.75, 1.75, 1.75, 1.75,
    1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75,
    1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75,
                1.75, 1.75, 1.75, 1.75, 1.75, 1.75,
    1.75, 1.75,
    1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75, 1.75]) / 0.52917721092)

atom_expected_radii = [1.0,  # Ghost atom
    1.0,  # 1
    0.927272,  # 2
    3.873661,  # 3
    2.849396,  # 4
    2.204757,  # 5
    1.714495,  # 6
    1.409631,  # 7
    1.232198,  # 8
    1.084786,  # 9
    0.965273,  # 10
    4.208762,  # 11
    3.252938,  # 12
    3.433889,  # 13
    2.752216,  # 14
    2.322712,  # 15
    2.060717,  # 16
    1.842024,  # 17
    1.662954,  # 18
    5.243652,  # 19
    4.218469,  # 20
    3.959716,  # 21
    3.778855,  # 22
    3.626288,  # 23
    3.675012,  # 24
    3.381917,  # 25
    3.258487,  # 26
    3.153572,  # 27
    3.059109,  # 28
    3.330979,  # 29
    2.897648,  # 30
    3.424103,  # 31
    2.866859,  # 32
    2.512233,  # 33
    2.299617,  # 34
    2.111601,  # 35
    1.951590,  # 36
    5.631401,  # 37
    4.632850,  # 38
    4.299870,  # 39
    4.091705,  # 40
    3.985219,  # 41
    3.841740,  # 42
    3.684647,  # 43
    3.735235,  # 44
    3.702057,  # 45
    1.533028,  # 46
    3.655961,  # 47
    3.237216,  # 48
    3.777242,  # 49
    3.248093,  # 50
    2.901067,  # 51
    2.691328,  # 52
    2.501704,  # 53
    2.337950]  # 54

def get_atomz(elmt: Union[str, ZType]) -> ZType:
    # returns the atomic number from the given element
    if isinstance(elmt, str):
        return periodic_table_atomz[elmt]
    elif isinstance(elmt, torch.Tensor):
        assert elmt.numel() == 1
        return elmt
    else:  # float or int
        return elmt

def get_atom_mass(atomz: int) -> float:
    # returns the atomic mass in atomic unit
    return atom_masses[atomz] * 1822.888486209

def get_period(atz: int) -> int:
    # get the period of the given atom z
    if atz <= 2:
        return 1
    elif atz <= 10:
        return 2
    elif atz <= 18:
        return 3
    elif atz <= 36:
        return 4
    elif atz <= 54:
        return 5
    elif atz <= 86:
        return 6
    elif atz <= 118:
        return 7
    else:
        raise RuntimeError("Unimplemented atomz: %d" % atz)


# 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 warnings

import numpy as np

from qutip.operators import tensor, identity, destroy, sigmax, sigmaz
from qutip.states import basis
from qutip.qip.circuit import QubitCircuit, Gate
from qutip.qip.device.processor import Processor
from qutip.qip.device.modelprocessor import ModelProcessor, GateDecomposer
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo

__all__ = ['DispersivecQED', 'CQEDGateDecomposer']

class DispersivecQED(ModelProcessor):
    """
    The processor based on the physical implementation of
    a dispersive cavity QED system.
    The available Hamiltonian of the system is predefined.
    For a given pulse amplitude matrix, the processor can
    calculate the state evolution under the given control pulse,
    either analytically or numerically.

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

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

    num_levels: int, optional
        The number of energy levels in the resonator.

    deltamax: int or list, optional
        The sigma-x paraicient for each of the qubits in the system.

    epsmax: int or list, optional
        The sigma-z paraicient for each of the qubits in the system.

    w0: int, optional
        The base frequency of the resonator.

    eps: int or list, optional
        The epsilon for each of the qubits in the system.

    delta: int or list, optional
        The epsilon for each of the qubits in the system.

    g: int or list, optional
        The interaction strength for each of the qubit with the resonator.

    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 value is a
        qubit system of ``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_ops: list
        A list of sigmax Hamiltonians for each qubit.

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

    cavityqubit_ops: list
        A list of interacting Hamiltonians between cavity and each qubit.

    sx_u: array_like
        Pulse matrix for sigmax Hamiltonians.

    sz_u: array_like
        Pulse matrix for sigmaz Hamiltonians.

    g_u: array_like
        Pulse matrix for interacting Hamiltonians
        between cavity and each qubit.

    wq: list of float
        The frequency of the qubits calculated from
        eps and delta for each qubit.

    Delta: list of float
        The detuning with repect to w0 calculated
        from wq and w0 for each qubit.
    """

    def __init__(self, N, correct_global_phase=True,
                 num_levels=10, deltamax=1.0,
                 epsmax=9.5, w0=10., wq=None, eps=9.5,
                 delta=0.0, g=0.01, t1=None, t2=None):
        super(DispersivecQED, 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"
        self.num_levels = num_levels
        self._paras = {}
        self.set_up_params(
            N=N, num_levels=num_levels, deltamax=deltamax,
            epsmax=epsmax, w0=w0, wq=wq, eps=eps,
            delta=delta, g=g)
        self.set_up_ops(N)
        self.dims = [num_levels] + [2] * N

    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.
        """
        # single qubit terms
        self.a = tensor([destroy(self.num_levels)] +
                        [identity(2) for n in range(N)])
        self.ctrls.append(self.a.dag() * self.a)
        self.ctrls += [tensor([identity(self.num_levels)] +
                              [sigmax() if m == n else identity(2)
                               for n in range(N)])
                       for m in range(N)]
        self.ctrls += [tensor([identity(self.num_levels)] +
                              [sigmaz() if m == n else identity(2)
                               for n in range(N)])
                       for m in range(N)]
        # interaction terms
        for n in range(N):
            sm = tensor([identity(self.num_levels)] +
                        [destroy(2) if m == n else identity(2)
                         for m in range(N)])
            self.ctrls.append(self.a.dag() * sm + self.a * sm.dag())

        self.psi_proj = tensor([basis(self.num_levels, 0)] +
                               [identity(2) for n in range(N)])

    def set_up_params(
            self, N, num_levels, deltamax,
            epsmax, w0, wq, eps, delta, g):
        """
        Save the parameters in the attribute `params` and check the validity.

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

        num_levels: int
            The number of energy levels in the resonator.

        deltamax: list
            The sigma-x paraicient for each of the qubits in the system.

        epsmax: list
            The sigma-z paraicient for each of the qubits in the system.

        wo: int
            The base frequency of the resonator.

        wq: list
            The frequency of the qubits.

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

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

        g: list
            The interaction strength for each of the qubit with the resonator.

        Notes
        -----
        All parameters will be multiplied by 2*pi for simplicity
        """
        sx_para = super(DispersivecQED, self)._para_list(deltamax, N)
        self._paras["sx"] = sx_para
        sz_para = super(DispersivecQED, self)._para_list(epsmax, N)
        self._paras["sz"] = sz_para
        w0 = w0 * 2 * np.pi
        self._paras["w0"] = w0
        eps = super(DispersivecQED, self)._para_list(eps, N)
        self._paras["eps"] = eps
        delta = super(DispersivecQED, self)._para_list(delta, N)
        self._paras["delta"] = delta
        g = super(DispersivecQED, self)._para_list(g, N)
        self._paras["g"] = g

        # computed
        self.wq = [np.sqrt(eps[i]**2 + delta[i]**2) for i in range(N)]
        self.Delta = [self.wq[i] - w0 for i in range(N)]

        # rwa/dispersive regime tests
        if any([g[i] / (w0 - self.wq[i]) > 0.05 for i in range(N)]):
            warnings.warn("Not in the dispersive regime")

        if any([(w0 - self.wq[i])/(w0 + self.wq[i]) > 0.05 for i in range(N)]):
            warnings.warn(
                "The rotating-wave approximation might not be valid.")

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

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

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

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

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

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

    def get_ops_labels(self):
        """
        Get the labels for each Hamiltonian.
        """
        return ([r"$a^\dagger a$"] +
                [r"$\sigma_x^%d$" % n for n in range(self.N)] +
                [r"$\sigma_z^%d$" % n for n in range(self.N)] +
                [r"$g_{%d}$" % (n) for n in range(self.N)])

    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.qc0.resolve_gates(
            basis=["SQRTISWAP", "ISWAP", "RX", "RZ"])
        return self.qc1

    def eliminate_auxillary_modes(self, U):
        """
        Eliminate the auxillary modes like the cavity modes in cqed.
        """
        return self.psi_proj.dag() * U * self.psi_proj

    def load_circuit(self, qc):
        """
        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.

        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 = CQEDGateDecomposer(
            self.N, self._paras, self.wq, self.Delta,
            global_phase=0., num_ops=len(self.ctrls))
        self.tlist, self.coeffs, self.global_phase = dec.decompose(gates)

        # TODO The amplitude of the first control a.dag()*a
        # was set to zero before I made this refactoring.
        # It is probably due to the fact that
        # it contributes only a constant (N) and can be neglected.
        # but change the below line to np.ones leads to test error.
        self.coeffs[0] = self._paras["w0"] * np.zeros((self.sx_u.shape[1]))
        return self.tlist, self.coeffs

class CQEDGateDecomposer(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.

    wq: list of float
        The frequency of the qubits calculated from
        eps and delta for each qubit.

    Delta: list of float
        The detuning with repect to w0 calculated
        from wq and w0 for each qubit.

    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.
    """
    def __init__(self, N, params, wq, Delta, global_phase, num_ops):
        super(CQEDGateDecomposer, 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._sx_ind = list(range(1, N + 1))
        self._sz_ind = list(range(N + 1, 2*N + 1))
        self._g_ind = list(range(2*N + 1, 3*N + 1))
        self.wq = wq
        self.Delta = Delta
        self.global_phase = global_phase

    def decompose(self, gates):
        tlist, coeffs = super(CQEDGateDecomposer, 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 sqrtiswap_dec(self, gate):
        """
        Decomposer for the SQRTISWAP gate

        Notes
        -----
        This version of sqrtiswap_dec has very low fidelity, please use
        iswap
        """
        # FIXME This decomposition has poor behaviour
        pulse = np.zeros(self.num_ops)
        q1, q2 = gate.targets
        pulse[self._sz_ind[q1]] = self.wq[q1] - self.params["w0"]
        pulse[self._sz_ind[q2]] = self.wq[q2] - self.params["w0"]
        pulse[self._g_ind[q1]] = self.params["g"][q1]
        pulse[self._g_ind[q2]] = self.params["g"][q2]
        J = self.params["g"][q1] * self.params["g"][q2] * (
            1 / self.Delta[q1] + 1 / self.Delta[q2]) / 2
        t = (4 * np.pi / abs(J)) / 8
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

        # corrections
        gate1 = Gate("RZ", [q1], None, arg_value=-np.pi/4)
        self.rz_dec(gate1)
        gate2 = Gate("RZ", [q2], None, arg_value=-np.pi/4)
        self.rz_dec(gate2)
        gate3 = Gate("GLOBALPHASE", None, None, arg_value=-np.pi/4)
        self.globalphase_dec(gate3)

    def iswap_dec(self, gate):
        """
        Decomposer for the ISWAP gate
        """
        pulse = np.zeros(self.num_ops)
        q1, q2 = gate.targets
        pulse[self._sz_ind[q1]] = self.wq[q1] - self.params["w0"]
        pulse[self._sz_ind[q2]] = self.wq[q2] - self.params["w0"]
        pulse[self._g_ind[q1]] = self.params["g"][q1]
        pulse[self._g_ind[q2]] = self.params["g"][q2]
        J = self.params["g"][q1] * self.params["g"][q2] * (
            1 / self.Delta[q1] + 1 / self.Delta[q2]) / 2
        t = (4 * np.pi / abs(J)) / 4
        self.dt_list.append(t)
        self.coeff_list.append(pulse)

        # corrections
        gate1 = Gate("RZ", [q1], None, arg_value=-np.pi/2.)
        self.rz_dec(gate1)
        gate2 = Gate("RZ", [q2], None, arg_value=-np.pi/2)
        self.rz_dec(gate2)
        gate3 = Gate("GLOBALPHASE", None, None, arg_value=-np.pi/2)
        self.globalphase_dec(gate3)

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


from abc import ABC, abstractmethod

from sklearn.base import BaseEstimator, RegressorMixin, MultiOutputMixin
from sklearn.utils.validation import check_is_fitted
from sklearn.model_selection import check_cv

from ._solvers import RIDGE_SOLVERS
from ._random_search import GROUP_RIDGE_SOLVERS
from ._random_search import solve_ridge_cv_svd

from ..validation import check_array
from ..validation import _get_string_dtype
from ..backend import get_backend
from ..backend import force_cpu_backend
from ..scoring import r2_score
from ..scoring import r2_score_split

class _BaseRidge(ABC, MultiOutputMixin, RegressorMixin, BaseEstimator):
    """Base class for ridge estimators"""

    @property
    @classmethod
    @abstractmethod
    def ALL_SOLVERS(cls):
        ...

    def _call_solver(self, **direct_params):
        if self.solver not in self.ALL_SOLVERS:
            raise ValueError("Unknown solver=%r." % self.solver)

        function = self.ALL_SOLVERS[self.solver]
        solver_params = self.solver_params or {}

        # check duplicated parameters
        intersection = set(direct_params.keys()).intersection(
            set(solver_params.keys()))
        if intersection:
            raise ValueError(
                'Parameters %s should not be given in solver_params, since '
                'they are either fixed or have a direct parameter in %s.' %
                (intersection, self.__class__.__name__))

        return function(**direct_params, **solver_params)

    def _more_tags(self):
        return {'requires_y': True}

class Ridge(_BaseRidge):
    """Ridge regression.

    Solve the ridge regression::

        b* = argmin_b ||X @ b - Y||^2 + alpha ||b||^2.

    Parameters
    ----------
    alpha : float, or array of shape (n_targets, )
        L2 regularization parameter.

    fit_intercept : boolean
        Whether to fit an intercept.
        If False, X and Y must be zero-mean over samples.

    solver : str
        Algorithm used during the fit, in {"svd"}.

    solver_params : dict or None
        Additional parameters for the solver.
        See more details in the docstring of the function:
        ``Ridge.ALL_SOLVERS[solver]``

    force_cpu : bool
        If True, computations will be performed on CPU, ignoring the
        current backend. If False, use the current backend.

    Attributes
    ----------
    coef_ : array of shape (n_features) or (n_features, n_targets)
        Ridge coefficients.

    intercept_ : float or array of shape (n_targets, )
        Intercept. Only present if fit_intercept is True.

    n_features_in_ : int
        Number of features used during the fit.

    dtype_ : str
        Dtype of input data.

    Examples
    --------
    >>> from himalaya.ridge import Ridge
    >>> import numpy as np
    >>> n_samples, n_features, n_targets = 10, 5, 3
    >>> X = np.random.randn(n_samples, n_features)
    >>> Y = np.random.randn(n_samples, n_targets)
    >>> model = Ridge()
    >>> model.fit(X, Y)
    Ridge()
    """
    ALL_SOLVERS = RIDGE_SOLVERS

    def __init__(self, alpha=1, fit_intercept=False, solver="svd",
                 solver_params=None, force_cpu=False):
        self.alpha = alpha
        self.fit_intercept = fit_intercept
        self.solver = solver
        self.solver_params = solver_params
        self.force_cpu = force_cpu

    @force_cpu_backend
    def fit(self, X, y=None):
        """Fit the model.

        Parameters
        ----------
        X : array of shape (n_samples, n_features)
            Training data.

        y : array of shape (n_samples,) or (n_samples, n_targets)
            Target values.

        Returns
        -------
        self : returns an instance of self.
        """
        # backend = get_backend()
        X = check_array(X, ndim=2)
        self.dtype_ = _get_string_dtype(X)
        y = check_array(y, dtype=self.dtype_, ndim=[1, 2])
        if X.shape[0] != y.shape[0]:
            raise ValueError("Inconsistent number of samples.")

        self.n_features_in_ = X.shape[1]

        ravel = False
        if y.ndim == 1:
            y = y[:, None]
            ravel = True

        # ------------------ call the solver
        tmp = self._call_solver(X=X, Y=y, alpha=self.alpha,
                                fit_intercept=self.fit_intercept)
        if self.fit_intercept:
            self.coef_, self.intercept_ = tmp
        else:
            self.coef_ = tmp

        if ravel:
            self.coef_ = self.coef_[:, 0]
            if self.fit_intercept:
                self.intercept_ = self.intercept_[0]

        return self

    @force_cpu_backend
    def predict(self, X):
        """Predict using the model.

        Parameters
        ----------
        X : array of shape (n_samples_test, n_features)
            Testing features.

        Returns
        -------
        Y_hat : array of shape (n_samples,) or (n_samples, n_targets)
            Returns predicted values.
        """
        check_is_fitted(self)
        backend = get_backend()
        X = check_array(X, dtype=self.dtype_, ndim=2)
        if X.shape[1] != self.n_features_in_:
            raise ValueError(
                'Different number of features in X than during fit.')

        Y_hat = backend.to_cpu(X) @ backend.to_cpu(self.coef_)
        if self.fit_intercept:
            Y_hat += backend.to_cpu(self.intercept_)
        return Y_hat

    @force_cpu_backend
    def score(self, X, y):
        """Return the coefficient of determination R^2 of the prediction.

        Parameters
        ----------
        X : array of shape (n_samples_test, n_features)
            Testing features.

        y : array-like of shape (n_samples,) or (n_samples, n_targets)
            True values for X.

        Returns
        -------
        score : array of shape (n_targets, )
            R^2 of self.predict(X) versus y.
        """
        y_pred = self.predict(X)
        y_true = check_array(y, dtype=self.dtype_, ndim=self.coef_.ndim)

        if y_true.ndim == 1:
            return r2_score(y_true[:, None], y_pred[:, None])[0]
        else:
            return r2_score(y_true, y_pred)

class RidgeCV(Ridge):
    """Ridge regression with efficient cross-validation over alpha.

    Solve the ridge regression::

        b* = argmin_b ||X @ b - Y||^2 + alpha ||b||^2,

    with a grid-search over cross-validation to find the best alpha.

    Parameters
    ----------
    alphas : array of shape (n_alphas, )
        List of L2 regularization parameter to try.

    fit_intercept : boolean
        Whether to fit an intercept.
        If False, X and Y must be zero-mean over samples.

    solver : str
        Algorithm used during the fit, "svd" only for now.

    solver_params : dict or None
        Additional parameters for the solver.
        See more details in the docstring of the function:
        ``RidgeCV.ALL_SOLVERS[solver]``

    cv : int or scikit-learn splitter
        Cross-validation splitter. If an int, KFold is used.

    Y_in_cpu : bool
        If True, keep the target values ``y`` in CPU memory (slower).

    force_cpu : bool
        If True, computations will be performed on CPU, ignoring the
        current backend. If False, use the current backend.

    Attributes
    ----------
    coef_ : array of shape (n_features) or (n_features, n_targets)
        Ridge coefficients.

    intercept_ : float or array of shape (n_targets, )
        Intercept. Only returned when fit_intercept is True.

    best_alphas_ : array of shape (n_targets, )
        Selected best hyperparameter alphas.

    cv_scores_ : array of shape (n_targets, )
        Cross-validation scores averaged over splits, for the best alpha.

    n_features_in_ : int
        Number of features used during the fit.

    Examples
    --------
    >>> from himalaya.ridge import RidgeCV
    >>> import numpy as np
    >>> n_samples, n_features, n_targets = 10, 5, 3
    >>> X = np.random.randn(n_samples, n_features)
    >>> Y = np.random.randn(n_samples, n_targets)
    >>> clf = RidgeCV()
    >>> clf.fit(X, Y)
    RidgeCV()
    """
    ALL_SOLVERS = dict(svd=solve_ridge_cv_svd)

    def __init__(self, alphas=[0.1, 1], fit_intercept=False, solver="svd",
                 solver_params=None, cv=5, Y_in_cpu=False, force_cpu=False):
        self.alphas = alphas
        self.fit_intercept = fit_intercept
        self.solver = solver
        self.solver_params = solver_params
        self.cv = cv
        self.Y_in_cpu = Y_in_cpu
        self.force_cpu = force_cpu

    @force_cpu_backend
    def fit(self, X, y=None):
        """Fit ridge regression model

        Parameters
        ----------
        X : array of shape (n_samples, n_features).
            Training data.

        y : array of shape (n_samples,) or (n_samples, n_targets)
            Target values.

        Returns
        -------
        self : returns an instance of self.
        """
        # backend = get_backend()
        X = check_array(X, ndim=2)
        self.dtype_ = _get_string_dtype(X)
        device = "cpu" if self.Y_in_cpu else None
        y = check_array(y, dtype=self.dtype_, ndim=[1, 2], device=device)
        if X.shape[0] != y.shape[0]:
            raise ValueError("Inconsistent number of samples.")

        alphas = check_array(self.alphas, dtype=self.dtype_, ndim=1)
        self.n_features_in_ = X.shape[1]

        ravel = False
        if y.ndim == 1:
            y = y[:, None]
            ravel = True

        cv = check_cv(self.cv)

        # ------------------ call the solver
        tmp = self._call_solver(X=X, Y=y, cv=cv, alphas=alphas,
                                fit_intercept=self.fit_intercept,
                                Y_in_cpu=self.Y_in_cpu)
        if self.fit_intercept:
            self.best_alphas_, self.coef_, self.cv_scores_ = tmp[:3]
            self.intercept_, = tmp[3:]
        else:
            self.best_alphas_, self.coef_, self.cv_scores_ = tmp

        self.cv_scores_ = self.cv_scores_[0]

        if ravel:
            self.coef_ = self.coef_[:, 0]
            if self.fit_intercept:
                self.intercept_ = self.intercept_[0]

        return self

###############################################################################
###############################################################################
###############################################################################
###############################################################################

class GroupRidgeCV(_BaseRidge):
    """Group ridge regression with cross-validation.

    Solve the group-regularized ridge regression::

        b* = argmin_b ||Z @ b - Y||^2 + ||b||^2

    where the feature space X_i is scaled by a group scaling ::

        Z_i = exp(deltas[i] / 2) X_i

    The solver optimizes the log group scalings ``deltas`` over
    cross-validation, using random search (``solver="random_search"``).

    Parameters
    ----------
    groups : array of shape (n_features, ), "input", or None
        Encoding of the group of each feature. If None, all features are
        gathered in one group, and the problem is equivalent to RidgeCV.
        If "input", the input features ``X`` should be a list of 2D arrays,
        corresponding to each group.

    solver : str
        Algorithm used during the fit, only "random_search" for now.

    solver_params : dict or None
        Additional parameters for the solver.
        See more details in the docstring of the function:
        ``GroupRidgeCV.ALL_SOLVERS[solver]``

    fit_intercept : boolean
        Whether to fit an intercept.
        If False, X and Y must be zero-mean over samples.

    cv : int or scikit-learn splitter
        Cross-validation splitter. If an int, KFold is used.

    random_state : int, or None
        Random generator seed. Use an int for deterministic search.

    Y_in_cpu : bool
        If True, keep the target values ``y`` in CPU memory (slower).

    force_cpu : bool
        If True, computations will be performed on CPU, ignoring the
        current backend. If False, use the current backend.

    Attributes
    ----------
    coef_ : array of shape (n_features) or (n_features, n_targets)
        Ridge coefficients.

    intercept_ : float or array of shape (n_targets, )
        Intercept. Only returned when fit_intercept is True.

    deltas_ : array of shape (n_groups, n_targets)
        Log of the group scalings.

    cv_scores_ : array of shape (n_iter, n_targets)
        Cross-validation scores, averaged over splits.

    n_features_in_ : int
        Number of features used during the fit.

    dtype_ : str
        Dtype of input data.

    best_alphas_ : array of shape (n_targets, )
        Equal to ``1. / exp(self.deltas_).sum(0)``. For the "random_search"
        solver, it corresponds to the best hyperparameter alphas, assuming that
        each squared group scaling vector sums to one (in particular, it is the
        case when ``solver_params['n_iter']`` is an integer).

    Examples
    --------
    >>> from himalaya.ridge import GroupRidgeCV
    >>> from himalaya.ridge import ColumnTransformerNoStack
    >>> from sklearn.pipeline import make_pipeline

    >>> # create a dataset
    >>> import numpy as np
    >>> n_samples, n_features, n_targets = 10, 5, 3
    >>> X = np.random.randn(n_samples, n_features)
    >>> Y = np.random.randn(n_samples, n_targets)

    >>> # Separate the first three columns and the last two
    >>> # columns, creating two groups of shape (n_samples, n_feature_i).
    >>> from sklearn.preprocessing import StandardScaler
    >>> ct = ColumnTransformerNoStack(
    ...     [("group_1", StandardScaler(), [0, 1, 2]),
    ...      ("group_2", StandardScaler(), slice(3, 5))])

    >>> # A model with automatic groups, as output by ColumnTransformerNoStack
    >>> model = GroupRidgeCV(groups="input")
    >>> pipe = make_pipeline(ct, model)
    >>> _ = pipe.fit(X, Y)
    """
    ALL_SOLVERS = GROUP_RIDGE_SOLVERS

    def __init__(self, groups=None, solver="random_search", solver_params=None,
                 fit_intercept=False, cv=5, random_state=None, Y_in_cpu=False,
                 force_cpu=False):

        self.groups = groups
        self.solver = solver
        self.solver_params = solver_params
        self.fit_intercept = fit_intercept
        self.cv = cv
        self.random_state = random_state
        self.Y_in_cpu = Y_in_cpu
        self.force_cpu = force_cpu

    @force_cpu_backend
    def fit(self, X, y=None):
        """Fit the model.

        Parameters
        ----------
        X : array of shape (n_samples, n_features), or list of length \
                (n_groups) with arrays of shape (n_samples, n_features)
            Training data.
            Must be a 2D array if ``groups`` is given.
            Must be a list of 2D arrays if ``groups="input"``.

        y : array of shape (n_samples,) or (n_samples, n_targets)
            Target values.

        Returns
        -------
        self : returns an instance of self.
        """
        backend = get_backend()

        Xs = self._split_groups(X, check=True)
        del X

        self.n_features_in_ = sum(Xi.shape[1] for Xi in Xs)

        self.dtype_ = _get_string_dtype(Xs[0])
        device = "cpu" if self.Y_in_cpu else None
        y = check_array(y, dtype=self.dtype_, ndim=[1, 2], device=device)

        if any([Xi.shape[0] != y.shape[0] for Xi in Xs]):
            raise ValueError("Inconsistent number of samples.")

        ravel = False
        if y.ndim == 1:
            y = y[:, None]
            ravel = True

        cv = check_cv(self.cv)

        # ------------------ call the solver
        tmp = self._call_solver(Xs=Xs, Y=y, cv=cv, return_weights=True,
                                random_state=self.random_state,
                                fit_intercept=self.fit_intercept,
                                Y_in_cpu=self.Y_in_cpu)
        if self.fit_intercept:
            self.deltas_, self.coef_, self.cv_scores_ = tmp[:3]
            self.intercept_, = tmp[3:]
        else:
            self.deltas_, self.coef_, self.cv_scores_ = tmp

        if self.solver == "random_search":
            self.best_alphas_ = 1. / backend.exp(self.deltas_).sum(0)
        else:
            self.best_alphas_ = None

        if ravel:
            self.coef_ = self.coef_[:, 0]
            self.deltas_ = self.deltas_[:, 0]
            if self.fit_intercept:
                self.intercept_ = self.intercept_[0]

        return self

    @force_cpu_backend
    def predict(self, X, split=False):
        """Predict using the model.

        Parameters
        ----------
        X : array of shape (n_samples_test, n_features), or list of length \
                (n_groups) with arrays of shape (n_samples_test, n_features)
            Training data.
            Must be a 2D array if ``groups`` is given.
            Must be a list of 2D arrays if ``groups="input"``.

        split : bool
            If True, the prediction is split on each feature space, and this
            method returns an array with one extra dimension (first dimension).
            The sum over this extra dimension corresponds to split=False.

        Returns
        -------
        Y_hat : array of shape (n_samples,) or (n_samples, n_targets)
            Returns predicted values.
            If parameter split is True, the array is of shape
            (n_groups, n_samples,) or (n_groups, n_samples, n_targets).
        """
        backend = get_backend()
        check_is_fitted(self)

        Xs = self._split_groups(X, dtype=self.dtype_, check=True)

        n_features = sum(Xi.shape[1] for Xi in Xs)
        if n_features != self.n_features_in_:
            raise ValueError(
                'Different number of features in X than during fit.')
        if split:
            if self.fit_intercept:
                raise NotImplementedError(
                    "Splitting the predictions is not implemented with "
                    "fit_intercept=True.")

            start = 0
            Ys_hat = None
            for ii, X_i in enumerate(Xs):
                n_features_i = X_i.shape[1]
                coef = self.coef_[start:start + n_features_i]
                start += n_features_i
                Y_hat = backend.to_cpu(X_i) @ backend.to_cpu(coef)
                if Ys_hat is None:
                    Ys_hat = backend.zeros_like(Y_hat,
                                                shape=(len(Xs), *Y_hat.shape))
                Ys_hat[ii] = Y_hat
            return Ys_hat

        else:
            X = backend.to_cpu(backend.concatenate(Xs, 1))
            del Xs

            Y_hat = X @ backend.to_cpu(self.coef_)
            if self.fit_intercept:
                Y_hat += backend.to_cpu(self.intercept_)
            return Y_hat

    @force_cpu_backend
    def score(self, X, y, split=False):
        """Return the coefficient of determination R^2 of the prediction.

        Parameters
        ----------
        X : array of shape (n_samples_test, n_features), or list of length \
                (n_groups) with arrays of shape (n_samples_test, n_features)
            Training data.
            Must be a 2D array if ``groups`` is given.
            Must be a list of 2D arrays if ``groups="input"``.

        y : array-like of shape (n_samples,) or (n_samples, n_targets)
            True values for X.

        split : bool
            If True, the prediction is split on each kernel, and the R2 score
            is decomposed over sub-predictions, adding an extra dimension
            in the first axis. The sum over this extra dimension corresponds to
            split=False.

        Returns
        -------
        score : array of shape (n_targets, ) or (n_kernels, n_targets)
            R^2 of self.predict(X) versus y.
            If parameter split is True, the array is of shape
            (n_kernels, n_targets).
        """
        y_pred = self.predict(X, split=split)
        y_true = check_array(y, dtype=self.dtype_, ndim=self.coef_.ndim)

        score_func = r2_score_split if split else r2_score

        if y_true.ndim == 1:
            return score_func(y_true[:, None], y_pred[:, None])[0]
        else:
            return score_func(y_true, y_pred)

    def _split_groups(self, X, check=True, **check_kwargs):
        backend = get_backend()

        # groups defined in X
        if isinstance(self.groups, str) and self.groups == "input":
            if check:
                X = [check_array(Xi, ndim=2, **check_kwargs) for Xi in X]
            return X

        # groups defined in self.groups
        if check:
            X = check_array(X, ndim=2, **check_kwargs)
        if self.groups is None:
            groups = backend.zeros((X.shape[1]))
        else:
            groups = self.groups
        groups = backend.asarray(groups)[:]
        Xs = [X[:, groups == u] for u in backend.unique(groups) if u >= 0]
        return Xs


"""
    Implements PPO

    PPO: https://arxiv.org/abs/1707.06347
    Modified from policy Written by Patrick Coady (pat-coady.github.io) to implement
    latest version of PPO with pessimistic ratio clipping

    o Has an option to servo both the learning rate and the clip_param to keep KL
      within  a specified range. This helps on some control tasks
      (i.e., Mujoco Humanid-v2)

    o Uses approximate KL

    o Models distribution of actions as a Gaussian with variance not conditioned on state

    o Has option to discretize sampled actions

"""
import numpy as np
import tensorflow as tf
from utils import Action_converter
import utils
import sklearn.utils

class Policy(object):
    """ NN-based policy approximation """
    def __init__(self, obs_dim, act_dim, actions_per_dim=3, kl_targ=0.003,epochs=20,discretize=False,
                 input_network_scale=10,output_network_scale=10,test_mode=False,shuffle=True,
                 entropy_coeff=0.0,eta=0.0, servo_kl=False, beta=0.1, lr=None):
        """
        Args:
            obs_dim:                num observation dimensions (int)
            act_dim:                num action dimensions (int)
            actions_per_dim:        used when discretizing action space
            kl_targ:                target KL divergence between pi_old and pi_new
            epochs:                 number of epochs per update
            discretize:             boolean, True discretizes action space
            input_network_scale:    NN input layer is of dim <input_network_scale> * obs_dim
            output_network_scale:   NN layer prior to output layer is of dim <input_network_scale> * act_dim
            test_mode:              boolean, True removes all exploration noise
            shuffle:                boolean, shuffles data each epoch
            entropy_coeff:          adds a loss term that encourages exploration.  This almost always makes things
                                    worse for control tasks
            eta:                    KL hingle loss parameter
            servo_kl:               boolean:  set to False to not servo beta to KL, which is original PPO implementation
            beta:                   clipping parameter for pessimistic loss ratio

        """
        print('PPO Policy 1')
        self.servo_kl = servo_kl
        self.input_network_scale = input_network_scale
        self.output_network_scale = output_network_scale
        self.test_mode = test_mode
        self.entropy_coeff = entropy_coeff
        self.discretize = discretize
        self.shuffle = shuffle
        self.actions_per_dim = actions_per_dim
        self.kl_stat = None
        self.entropy_stat = None
        self.eta = eta  # multiplier for D_KL-kl_targ hinge-squared loss
        self.kl_targ = kl_targ
        self.epochs = epochs
        self.lr = lr
        self.lr_multiplier = 1.0  # dynamically adjust lr when D_KL out of control
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.max_beta = 0.5
        self.min_beta = 0.01
        self.beta = beta
        self._build_graph()
        self._init_session()
        self.action_converter = Action_converter(1,actions_per_dim)
        print(self.input_network_scale)
        print('Actor Test Mode: ',self.test_mode)
        print('clip param: ',self.beta)

    def _build_graph(self):
        """ Build and initialize TensorFlow graph """
        self.g = tf.Graph()
        with self.g.as_default():
            self._placeholders()
            self._policy_nn()
            self._logprob()
            self._kl_entropy()
            self._sample()
            self._loss_train_op()
            self.init = tf.global_variables_initializer()

    def _placeholders(self):
        """ Input placeholders"""
        # observations, actions and advantages:
        self.obs_ph = tf.placeholder(tf.float32, (None, self.obs_dim), 'obs')
        self.act_ph = tf.placeholder(tf.float32, (None, self.act_dim), 'act')
        self.advantages_ph = tf.placeholder(tf.float32, (None,), 'advantages')
        # strength of D_KL loss terms:
        self.beta_ph = tf.placeholder(tf.float32, (), 'beta')
        self.eta_ph = tf.placeholder(tf.float32, (), 'eta')
        # learning rate:
        self.lr_ph = tf.placeholder(tf.float32, (), 'eta')
        # log_vars and means with pi_old (previous step's policy parameters):
        self.old_logp_ph =  tf.placeholder(tf.float32, (None,), 'old_logp')

        #self.old_log_vars_ph = tf.placeholder(tf.float32, (self.act_dim,), 'old_log_vars')
        #self.old_means_ph = tf.placeholder(tf.float32, (None, self.act_dim), 'old_means')

    def _policy_nn(self):
        """ Neural net for policy approximation function

        Policy parameterized by Gaussian means and variances. NN outputs mean
         action based on observation. Trainable variables hold log-variances
         for each action dimension (i.e. variances not determined by NN).
        """
        # hidden layer sizes determined by obs_dim and act_dim (hid2 is geometric mean)
        hid1_size = self.obs_dim * self.input_network_scale  # 10 empirically determined
        hid3_size = self.act_dim * self.output_network_scale  # 10 empirically determined
        hid2_size = int(np.sqrt(hid1_size * hid3_size))
        # heuristic to set learning rate based on NN size (tuned on 'Hopper-v1')
        if self.lr is None:
            self.lr = 9e-4 / np.sqrt(hid2_size)  # 9e-4 empirically determined
        # 3 hidden layers with tanh activations
        out = tf.layers.dense(self.obs_ph, hid1_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / self.obs_dim)), name="h1")
        out = tf.layers.dense(out, hid2_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / hid1_size)), name="h2")
        out = tf.layers.dense(out, hid3_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / hid2_size)), name="h3")
        self.means = tf.layers.dense(out, self.act_dim,
                                     kernel_initializer=tf.random_normal_initializer(
                                         stddev=np.sqrt(1 / hid3_size)), name="means")
        # logvar_speed is used to 'fool' gradient descent into making faster updates
        # to log-variances. heuristic sets logvar_speed based on network size.
        logvar_speed = (10 * hid3_size) // 48
        log_vars = tf.get_variable('logvars', (logvar_speed, self.act_dim), tf.float32,
                                   tf.constant_initializer(0.0))
        self.log_vars = tf.reduce_sum(log_vars, axis=0) - 1.0

        print('Policy Params -- h1: {}, h2: {}, h3: {}, lr: {:.3g}, logvar_speed: {}'
              .format(hid1_size, hid2_size, hid3_size, self.lr, logvar_speed))

    def _policy_rnn(self):
        """ Neural net for policy approximation function

        Policy parameterized by Gaussian means and variances. NN outputs mean
         action based on observation. Trainable variables hold log-variances
         for each action dimension (i.e. variances not determined by NN).
        """
        # hidden layer sizes determined by obs_dim and act_dim (hid2 is geometric mean)
        hid1_size = self.obs_dim * self.input_network_scale  # 10 empirically determined
        hid3_size = self.act_dim * self.output_network_scale  # 10 empirically determined
        hid2_size = int(np.sqrt(hid1_size * hid3_size))
        # heuristic to set learning rate based on NN size (tuned on 'Hopper-v1')
        if self.lr is None:
            self.lr = 9e-4 / np.sqrt(hid2_size)  # 9e-4 empirically determined
        # 3 hidden layers with tanh activations
        out1 = tf.layers.dense(self.obs_ph, hid1_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / self.obs_dim)), name="h1")
        out2 = tf.layers.dense(tf.concat(concat_dim=1,values=[out1,out3]), hid2_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / hid1_size)), name="h2")
        out3 = tf.layers.dense(out2, hid3_size, tf.tanh,
                              kernel_initializer=tf.random_normal_initializer(
                                  stddev=np.sqrt(1 / hid2_size)), name="h3")
        self.means = tf.layers.dense(out3, self.act_dim,
                                     kernel_initializer=tf.random_normal_initializer(
                                         stddev=np.sqrt(1 / hid3_size)), name="means")
        # logvar_speed is used to 'fool' gradient descent into making faster updates
        # to log-variances. heuristic sets logvar_speed based on network size.
        logvar_speed = (10 * hid3_size) // 48
        log_vars = tf.get_variable('logvars', (logvar_speed, self.act_dim), tf.float32,
                                   tf.constant_initializer(0.0))
        self.log_vars = tf.reduce_sum(log_vars, axis=0) - 1.0

        print('Policy Params -- h1: {}, h2: {}, h3: {}, lr: {:.3g}, logvar_speed: {}'
              .format(hid1_size, hid2_size, hid3_size, self.lr, logvar_speed))

    def _logprob(self):
        """ Calculate log probabilities of a batch of observations & actions

        Calculates log probabilities using previous step's model parameters and
        new parameters being trained.
        """
        logp = -0.5 * tf.reduce_sum(self.log_vars)
        logp += -0.5 * tf.reduce_sum(tf.square(self.act_ph - self.means) /
                                     tf.exp(self.log_vars), axis=1)
        logp += -0.5 * np.log(2.0 * np.pi) * self.act_dim
        self.logp = logp

    def _kl_entropy(self):
        """
        Add to Graph:
            1. KL divergence between old and new distributions
            2. Entropy of present policy given states and actions

        https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback.E2.80.93Leibler_divergence
        https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Entropy
        """

        self.kl = 0.5 * tf.reduce_mean(tf.square(self.logp - self.old_logp_ph))

        self.entropy = 0.5 * (self.act_dim * (np.log(2 * np.pi) + 1) +
                              tf.reduce_sum(self.log_vars))

    def _sample(self):
        """ Sample from distribution, given observation """
        self.sampled_act = (self.means +
                            tf.exp(self.log_vars / 2.0) *
                            tf.random_normal(shape=(self.act_dim,)))

        self.deterministic_act = (self.means)

    def _loss_train_op(self):
        ratio = tf.exp(self.logp - self.old_logp_ph)
        surr1 = self.advantages_ph * ratio
        surr2 = self.advantages_ph * tf.clip_by_value(ratio, 1.0 - self.beta_ph, 1.0 + self.beta_ph)

        loss1 = -tf.reduce_mean(tf.minimum(surr1,surr2)) - self.entropy_coeff * self.entropy
        loss2 = self.eta_ph * tf.square(tf.maximum(0.0, self.kl - 2.0 * self.kl_targ))

        self.loss = loss1 + loss2
        optimizer = tf.train.AdamOptimizer(self.lr_ph)
        self.train_op = optimizer.minimize(self.loss)

    def train(self, observes, actions, advantages, old_logp_np):

        feed_dict = {self.obs_ph: observes,
                     self.act_ph: actions,
                     self.advantages_ph: advantages,
                     self.old_logp_ph: old_logp_np,
                     self.beta_ph: self.beta,
                     self.eta_ph: self.eta,
                     self.lr_ph: self.lr * self.lr_multiplier}

        _, loss, kl, entropy, log_var_monitor = self.sess.run([self.train_op, self.loss, self.kl, self.entropy, self.log_vars], feed_dict)

        return loss, kl, entropy, log_var_monitor

    def _init_session(self):
        """Launch TensorFlow session and initialize variables"""
        self.sess = tf.Session(graph=self.g)
        self.sess.run(self.init)

    def sample(self, obs):
        """Draw sample from policy distribution"""
        feed_dict = {self.obs_ph: obs}

        if self.test_mode:
            action = self.sess.run(self.deterministic_act , feed_dict=feed_dict)
        else:
            action = self.sess.run(self.sampled_act , feed_dict=feed_dict)

        if self.discretize:
            idx = self.action_converter.action2idx(action[0])
            discrete_action = self.action_converter.idx2action(idx)
            env_action = discrete_action
        else:
            env_action = action
        #print('env act: ',env_action)
        return action, env_action

    def update(self, observes, actions, advantages, logger):
        """ Update policy based on observations, actions and advantages

        Args:
            observes: observations, shape = (N, obs_dim)
            actions: actions, shape = (N, act_dim)
            advantages: advantages, shape = (N,)
            logger: Logger object, see utils.py
        """

        feed_dict = {self.obs_ph: observes,
                     self.act_ph: actions,
                     self.advantages_ph: advantages}
        old_logp_np = np.squeeze(self.sess.run([self.logp],feed_dict))

        loss, kl, entropy = 0, 0, 0

        for e in range(self.epochs):

            if self.shuffle:
                    observes, actions, advantages, old_logp_np = sklearn.utils.shuffle(observes,actions,advantages,old_logp_np)

            loss, kl, entropy, log_var_monitor = self.train(observes, actions, advantages, old_logp_np)

            if kl > 4.0 * self.kl_targ and self.servo_kl:
                print(' *** BROKE ***')
                break

        if self.servo_kl:
            self.adjust_beta(kl)
        print('kl = ',kl, ' beta = ',self.beta,' lr_mult = ',self.lr_multiplier)
        self.kl_stat = kl
        self.entropy_stat = entropy
        var_monitor = np.exp(log_var_monitor/2.0)
        print('var: ' ,var_monitor),
        logger.log({'PolicyLoss': loss,
                    'PolicyEntropy': entropy,
                    'KL': kl,
                    'Beta': self.beta,
                    'Variance' : np.max(var_monitor),
                    'lr_multiplier': self.lr_multiplier})

    def adjust_beta(self,kl):
        if  kl < self.kl_targ / 2:
            self.beta = np.minimum(self.max_beta, 1.5 * self.beta)  # max clip beta
            #print('too low')
            if self.beta > (self.max_beta/2) and self.lr_multiplier < 10:
                self.lr_multiplier *= 1.5
        elif kl > self.kl_targ * 2:
            #print('too high')
            self.beta = np.maximum(self.min_beta, self.beta / 1.5)  # min clip beta
            if self.beta <= (2*self.min_beta) and self.lr_multiplier > 0.1:
                self.lr_multiplier /= 1.5

    def close_sess(self):
        """ Close TensorFlow session """
        self.sess.close()


"""
This module includes the description of a (multiturn) WakeField as well
as the implementation of the WakeSource objects.

A WakeField is defined as a composition of the elementary WakeKick
objects (see .wake_kicks module). They originate from WakeSources,
e.g. a WakeTable, Resonator and/or a ResistiveWall. The WakeField does
not directly accept the WakeKick objects, but takes a list of
WakeSources first (can be of different kinds), each of which knows how
to generate its WakeKick objects via the factory method
WakeSource.get_wake_kicks(..). The collection of WakeKicks from all the
WakeSources define the WakeField and are the elementary objects that are
stored, (i.e. the WakeField forgets about the origin of the WakeKicks
once they have been created).

@author Hannes Bartosik, Kevin Li, Giovanni Rumolo, Michael Schenk
@date March 2014
@brief Implementation of a WakeField as a composition of WakeKicks
       originating from different WakeSources.
@copyright CERN
"""

import numpy as np
from collections import deque
from scipy.constants import c, physical_constants
from scipy.interpolate import interp1d
from abc import ABCMeta, abstractmethod

from PyHEADTAIL.impedances.wake_kicks import *
from PyHEADTAIL.general.element import Element, Printing
from PyHEADTAIL.general.decorators import deprecated

sin = np.sin
cos = np.cos

def check_wake_sampling(bunch, slicer, wakes, beta=1, wake_column=None, bins=False):
    '''
    Handy function for quick visual check of sampling of the wake functions.
    For now only implemented for wake table type wakes.
    '''
    from scipy.constants import c
    import matplotlib.pyplot as plt

    ss = bunch.get_slices(slicer).z_centers
    zz = bunch.get_slices(slicer).z_bins
    ss = ss[:-1]
    ll = bunch.get_slices(slicer).lambda_z(ss, sigma=100)
    # ss = np.concatenate((s.z_centers-s.z_centers[-1], (s.z_centers-s.z_centers[0])[1:]))

    A = [wakes.wake_table['time'] * beta*c*1e-9, wakes.wake_table[wake_column] * 1e15]
    W = [ss[::-1], wakes.function_transverse(wake_column)(beta, ss)]

    fig, (ax1, ax2) = plt.subplots(2, figsize=(16,12), sharex=True)

    ax1.plot(ss, ll)

    ax2.plot(A[0], (A[1]), 'b-+', ms=12)
    ax2.plot(W[0][:-1], (-1*W[1][1:]), 'r-x')
    if bins:
        [ax2.axvline(z, color='g') for z in zz]

    ax2.grid()
    lgd = ['Table', 'Interpolated']
    if bins:
        lgd += ['Bin edges']
    ax2.legend(lgd)

    print('\n--> Resulting number of slices: {:g}'.format(len(ss)))

    return ax1

class WakeField(Element):
    """ A WakeField is defined by elementary WakeKick objects that may
    originate from different WakeSource objects. Usually, there is
    no need for the user to define more than one instance of the
    WakeField class in a simulation - except if one wants to use
    different slicing configurations (one WakeField object is allowed
    to have exactly one slicing configuration, i.e. only one instance
    of the Slicer class). A WakeField also is able to calculate the wake
    forces coming from earlier turns (multiturn wakes) by archiving the
    longitudinal bunch distribution (SliceSet instances) a number of
    turns back. """

    def __init__(self, slicer, *wake_sources):
        """ Accepts a list of WakeSource objects. Each WakeSource object
        knows how to generate its corresponding WakeKick objects. The
        collection of all the WakeKick objects of each of the passed
        WakeSource objects defines the WakeField.
        When instantiating the WakeField object, the WakeKick objects
        for each WakeSource defined in wake_sources are requested. The
        returned WakeKick lists are all stored in the
        WakeField.wake_kicks list. The WakeField itself forgets about
        the origin (WakeSource) of the kicks as soon as they have been
        generated.
        Exactly one instance of the Slicer class must be passed to the
        WakeField constructor. All the wake field components (kicks)
        hence use the same slicing and thus the same slice_set to
        calculate the strength of the kicks.
        To calculate the contributions from multiturn wakes, the
        longitudinal beam distributions (SliceSet instances) are
        archived in a deque. In parallel to the slice_set_deque,
        there is a slice_set_age_deque to keep track of the age of
        each of the SliceSet instances."""
        self.slicer = slicer

        self.wake_kicks = []
        for source in wake_sources:
            kicks = source.get_wake_kicks(self.slicer)
            self.wake_kicks.extend(kicks)

        n_turns_wake_max = max([ source.n_turns_wake
                                 for source in wake_sources ])
        self.slice_set_deque = deque([], maxlen=n_turns_wake_max)
        self.slice_set_age_deque = deque([], maxlen=n_turns_wake_max)

    def track(self, bunch):
        """ Calls the WakeKick.apply(bunch, slice_set) method of each of
        the WakeKick objects stored in self.wake_kicks. A slice_set is
        necessary to perform this operation. It is requested from the
        bunch (instance of the Particles class) using the
        Particles.get_slices(self.slicer) method, where self.slicer is
        the instance of the Slicer class used for this particluar
        WakeField object. A slice_set is returned according to the
        self.slicer configuration. The statistics mean_x and mean_y are
        requested to be calculated and saved in the SliceSet instance,
        too, s.t. the first moments x, y can be calculated by the
        WakeKick instances. """

        # Update ages of stored SliceSet instances.
        for i in range(len(self.slice_set_age_deque)):
            self.slice_set_age_deque[i] += (
                bunch.circumference / (bunch.beta * c))

        slice_set = bunch.get_slices(self.slicer,
                                     statistics=['mean_x', 'mean_y'])
        self.slice_set_deque.appendleft(slice_set)
        self.slice_set_age_deque.appendleft(0.)

        for kick in self.wake_kicks:
            kick.apply(bunch, self.slice_set_deque, self.slice_set_age_deque)

''' WakeSource classes. '''

class WakeSource(Printing, metaclass=ABCMeta):
    """ Abstract base class for wake sources, such as WakeTable,
    Resonator or ResistiveWall. """

    @abstractmethod
    def get_wake_kicks(self, slicer_mode):
        """ Factory method. Creates instances of the WakeKick objects
        for the given WakeSource and returns them as a list wake_kicks.
        This method is usually only called by a WakeField object to
        collect and create all the WakeKick objects originating from the
        different sources. (The slicer mode Slicer.mode must be passed
        at instantiation of a WakeKick object only to set the
        appropriate convolution method. See docstrings of WakeKick
        class.) """
        pass

class WakeTable(WakeSource):
    """ Class to define wake functions and WakeKick objects using wake
    data from a table. """

    def __init__(self, wake_file, wake_file_columns, n_turns_wake=1,
                 *args, **kwargs):
        """ Load data from the wake_file and store them in a dictionary
        self.wake_table. Keys are the names specified by the user in
        wake_file_columns and describe the names of the wake field
        components (e.g. dipole_x or dipole_yx). The dict values are
        given by the corresponding data read from the table. The
        nomenclature of the wake components must be strictly obeyed.
        Valid names for wake components are:

        'constant_x', 'constant_y', 'dipole_x', 'dipole_y', 'dipole_xy',
        'dipole_yx', 'quadrupole_x', 'quadrupole_y', 'quadrupole_xy',
        'quadrupole_yx', 'longitudinal'.

        The order of wake_file_columns is relevant and must correspond
        to the one in the wake_file. There is no way to check this here
        and it is in the responsibility of the user to ensure it is
        correct. Two checks made here are whether the length of
        wake_file_columns corresponds to the number of columns in the
        wake_file and whether a column 'time' is specified.

        The units and signs of the wake table data are assumed to follow
        the HEADTAIL conventions, i.e.
          time: [ns]
          transverse wake components: [V/pC/mm]
          longitudinal wake component: [V/pC].

        The parameter 'n_turns_wake' defines how many turns are
        considered for the multiturn wakes. It is 1 by default, i.e.
        multiturn wakes are off. """
        super(WakeTable, self).__init__(*args, **kwargs)

        self.wake_table = {}

        wake_data = np.loadtxt(wake_file)
        if len(wake_file_columns) != wake_data.shape[1]:
            raise ValueError("Length of wake_file_columns list does not" +
                             " correspond to the number of columns in the" +
                             " specified wake_file. \n")
        if 'time' not in wake_file_columns:
            raise ValueError("No wake_file_column with name 'time' has" +
                             " been specified. \n")

        for i, column_name in enumerate(wake_file_columns):
            self.wake_table.update({ column_name : wake_data[:,i] })

        self.n_turns_wake = n_turns_wake

    def get_wake_kicks(self, slicer):
        """ Factory method. Creates instances of the appropriate
        WakeKick objects for all the wake components provided by the
        user (and the wake table data). The WakeKick objects are
        returned as a list wake_kicks. """
        wake_kicks = []

        # Constant wake kicks.
        if self._is_provided('constant_x'):
            wake_function = self.function_transverse('constant_x')
            wake_kicks.append(ConstantWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('constant_y'):
            wake_function = self.function_transverse('constant_y')
            wake_kicks.append(ConstantWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('longitudinal'):
            wake_function = self.function_longitudinal()
            wake_kicks.append(ConstantWakeKickZ(
                wake_function, slicer, self.n_turns_wake))

        # Dipolar wake kicks.
        if self._is_provided('dipole_x'):
            wake_function = self.function_transverse('dipole_x')
            wake_kicks.append(DipoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('dipole_y'):
            wake_function = self.function_transverse('dipole_y')
            wake_kicks.append(DipoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('dipole_xy'):
            wake_function = self.function_transverse('dipole_xy')
            wake_kicks.append(DipoleWakeKickXY(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('dipole_yx'):
            wake_function = self.function_transverse('dipole_yx')
            wake_kicks.append(DipoleWakeKickYX(
                wake_function, slicer, self.n_turns_wake))

        # Quadrupolar wake kicks.
        if self._is_provided('quadrupole_x'):
            wake_function = self.function_transverse('quadrupole_x')
            wake_kicks.append(QuadrupoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('quadrupole_y'):
            wake_function = self.function_transverse('quadrupole_y')
            wake_kicks.append(QuadrupoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('quadrupole_xy'):
            wake_function = self.function_transverse('quadrupole_xy')
            self.kicks.append(QuadrupoleWakeKickXY(
                wake_function, slicer, self.n_turns_wake))

        if self._is_provided('quadrupole_yx'):
            wake_function = self.function_transverse('quadrupole_yx')
            wake_kicks.append(QuadrupoleWakeKickYX(
                wake_function, slicer, self.n_turns_wake))

        return wake_kicks

    def _is_provided(self, wake_component):
        """ Check whether wake_component is a valid name and available
        in wake table data. Return 'True' if yes and 'False' if no. """
        if wake_component in list(self.wake_table.keys()):
            return True
        else:
            # self.warns(wake_component + ' \n' +
            #       'Wake component is either not provided or does not \n'+
            #       'use correct nomenclature. See docstring of WakeTable \n' +
            #       'constructor to display valid names. \n')
            return False

    def function_transverse(self, wake_component):
        """ Defines and returns the wake(beta, dz) function for the
        given wake_component (transverse). Data from the wake table are
        used, but first converted to SI units assuming that time is
        specified in [ns] and transverse wake field strengths in
        [V/pC/mm]. Sign conventions are applied (HEADTAIL conventions).
        dz is related to wake table time data by dz = beta c dt (dz < 0
        for the ultrarelativistic case).
        The wake(dt) uses the scipy.interpolate.interp1d linear
        interpolation to calculate the wake strength at an arbitrary
        value of dt (provided it is in the valid range). The valid range
        of dt is given by the time range from the wake table. If values
        of wake(dt) are requested for dt outside the valid range, a
        ValueError is raised by interp1d.
        Very basic conformity checks for the wake table data are already
        performed at definition time of the wake(dt) method. E.g.
        whether the specified wake is valid only for ultrarelativistic
        cases or low beta cases. In the former case, the wake strength
        at time 0 must be defined by the user! """
        convert_to_s = 1e-9
        convert_to_V_per_Cm = 1e15

        time = convert_to_s * self.wake_table['time']
        wake_strength = -convert_to_V_per_Cm * self.wake_table[wake_component]
        interpolation_function = interp1d(time, wake_strength)

        if (time[0] == 0) and (wake_strength[0] == 0):
            def wake(dt, *args, **kwargs):
                dt = dt.clip(max=0)
                return interpolation_function(-dt)
            self.prints(wake_component +
                  ' Assuming ultrarelativistic wake.')

        elif (time[0] < 0):
            def wake(dt, *args, **kwargs):
                return interpolation_function(-dt)
            self.prints(wake_component +  ' Found low beta wake.')

        else:
            raise ValueError(wake_component +
                             ' does not meet requirements.')
        return wake

    def function_longitudinal(self):
        """ Defines and returns the wake(dt) function for the given
        wake_component (longitudinal). Data from the wake table are
        used, but first converted to SI units assuming that time is
        specified in [ns] and longitudinal wake field strength in
        [V/pC]. Sign conventions are applied (HEADTAIL conventions).
        The wake(dt) uses the scipy.interpolate.interp1d linear
        interpolation to calculate the wake strength at an arbitrary
        value of dt (provided it is in the valid range). The valid range
        of dt is given by the time range from the wake table. If values
        of wake(dt) are requested for dt outside the valid range, a
        ValueError is raised by interp1d.
        The beam loading theorem is respected and applied for dt=0. """
        convert_to_s = 1e-9
        convert_to_V_per_C = 1e12

        time = convert_to_s * self.wake_table['time']
        wake_strength = -convert_to_V_per_C * self.wake_table['longitudinal']
        interpolation_function = interp1d(time, wake_strength)

        def wake(dt, *args, **kwargs):
            wake_interpolated = interpolation_function(-dt)
            if time[0] == 0:
                # Beam loading theorem: Half value of wake strength at
                # dt = 0.
                return (np.sign(-dt) + 1.) / 2. * wake_interpolated
            elif time[0] < 0:
                return wake_interpolated
            else:
                raise ValueError('Longitudinal wake component does not meet' +
                                 ' requirements.')
        return wake

class Resonator(WakeSource):
    """ Class to describe the wake functions originating from a
    resonator impedance. Alex Chao's resonator model (Eq. 2.82) is used
    as well as the definitions from HEADTAIL. """

    def __init__(self, R_shunt, frequency, Q,
                 Yokoya_X1, Yokoya_Y1, Yokoya_X2, Yokoya_Y2, switch_Z,
                 n_turns_wake=1, *args, **kwargs):
        """ General constructor to create a Resonator WakeSource object
        describing the wake functions of a resonator impedance. Alex
        Chao's resonator model (Eq. 2.82) is used as well as definitions
        from HEADTAIL.
        Note that it is no longer allowed to pass a LIST of parameters
        to generate a number of resonators with different parameters
        within the same Resonator object. Instead, create the Resonator
        objects and pass all of them to the WakeField constructor.
        The parameter 'n_turns_wake' defines how many turns are
        considered for the multiturn wakes. It is 1 by default, i.e.
        multiturn wakes are off. """
        super(Resonator, self).__init__(*args, **kwargs)

        self.R_shunt = R_shunt
        self.frequency = frequency
        self.Q = Q
        self.Yokoya_X1 = Yokoya_X1
        self.Yokoya_X2 = Yokoya_X2
        self.Yokoya_Y1 = Yokoya_Y1
        self.Yokoya_Y2 = Yokoya_Y2
        self.switch_Z = switch_Z
        self.n_turns_wake = n_turns_wake

    def get_wake_kicks(self, slicer):
        """ Factory method. Creates instances of the appropriate
        WakeKick objects for a Resonator WakeSource with the specified
        parameters. A WakeKick object is instantiated only if the
        corresponding Yokoya factor is non-zero. The WakeKick objects
        are returned as a list wake_kicks. """
        wake_kicks = []

        # Dipole wake kick x.
        if self.Yokoya_X1:
            wake_function = self.function_transverse(self.Yokoya_X1)
            wake_kicks.append(DipoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        # Quadrupole wake kick x.
        if self.Yokoya_X2:
            wake_function = self.function_transverse(self.Yokoya_X2)
            wake_kicks.append(QuadrupoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        # Dipole wake kick y.
        if self.Yokoya_Y1:
            wake_function = self.function_transverse(self.Yokoya_Y1)
            wake_kicks.append(DipoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        # Quadrupole wake kick y.
        if self.Yokoya_Y2:
            wake_function = self.function_transverse(self.Yokoya_Y2)
            wake_kicks.append(QuadrupoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        # Constant wake kick z.
        if self.switch_Z:
            wake_function = self.function_longitudinal()
            wake_kicks.append(ConstantWakeKickZ(
                wake_function, slicer, self.n_turns_wake))

        return wake_kicks

    def function_transverse(self, Yokoya_factor):
        """ Define the wake function (transverse) of a resonator with
        the given parameters according to Alex Chao's resonator model
        (Eq. 2.82) and definitions of the resonator in HEADTAIL. """
        omega = 2 * np.pi * self.frequency
        alpha = omega / (2 * self.Q)
        omegabar = np.sqrt(np.abs(omega**2 - alpha**2))

        def wake(dt, *args, **kwargs):
            dt = dt.clip(max=0)
            if self.Q > 0.5:
                y = (Yokoya_factor * self.R_shunt * omega**2 / (self.Q *
                     omegabar) * np.exp(alpha*dt) * sin(omegabar*dt))
            elif self.Q == 0.5:
                y = (Yokoya_factor * self.R_shunt * omega**2 / self.Q *
                      np.exp(alpha*dt) * dt)
            else:
                y = (Yokoya_factor * self.R_shunt * omega**2 / (self.Q *
                     omegabar) * np.exp(alpha*dt) * np.sinh(omegabar*dt))
            return y
        return wake

    def function_longitudinal(self):
        """ Define the wake function (longitudinal) of a resonator with
        the given parameters according to Alex Chao's resonator model
        (Eq. 2.82) and definitions of the resonator in HEADTAIL. """
        omega = 2 * np.pi * self.frequency
        alpha = omega / (2 * self.Q)
        omegabar = np.sqrt(np.abs(omega**2 - alpha**2))

        def wake(dt, *args, **kwargs):
            if self.Q > 0.5:
                y = (-(np.sign(dt) - 1) * self.R_shunt * alpha *
                     np.exp(alpha * dt) * (cos(omegabar * dt) +
                     alpha / omegabar * sin(omegabar*dt)))
            elif self.Q == 0.5:
                y = (-(np.sign(dt) - 1) * self.R_shunt * alpha *
                     np.exp(alpha * dt) * (1. + alpha * dt))
            elif self.Q < 0.5:
                y = (-(np.sign(dt) - 1) * self.R_shunt * alpha *
                     np.exp(alpha * dt) * (np.cosh(omegabar * dt) +
                     alpha / omegabar * np.sinh(omegabar * dt)))
            return y
        return wake

class CircularResonator(Resonator):
    '''Circular Resonator.'''
    def __init__(self, R_shunt, frequency, Q, n_turns_wake=1,
                 *args, **kwargs):
        """ Special case of circular resonator. """
        Yokoya_X1 = 1.
        Yokoya_Y1 = 1.
        Yokoya_X2 = 0.
        Yokoya_Y2 = 0.
        switch_Z  = False

        super(CircularResonator, self).__init__(
            R_shunt, frequency, Q, Yokoya_X1, Yokoya_Y1,
            Yokoya_X2, Yokoya_Y2, switch_Z, n_turns_wake, *args, **kwargs)

class ParallelHorizontalPlatesResonator(Resonator):
    '''Broad-band resonator for horizontal parallel plates.'''
    def __init__(self, R_shunt, frequency, Q, n_turns_wake=1,
                 *args, **kwargs):
        """ Special case of parallel plate resonator. """
        Yokoya_X1 = np.pi**2 / 24.
        Yokoya_Y1 = np.pi**2 / 12.
        Yokoya_X2 = -np.pi**2 / 24.
        Yokoya_Y2 = np.pi**2 / 24.
        switch_Z  = False

        super(ParallelHorizontalPlatesResonator, self).__init__(
            R_shunt, frequency, Q, Yokoya_X1, Yokoya_Y1,
            Yokoya_X2, Yokoya_Y2, switch_Z, n_turns_wake, *args, **kwargs)

@deprecated('--> "ParallelPlatesResonator" will be removed '
            'in the near future. '
            'Use "ParallelHorizontalPlatesResonator" instead.\n')
class ParallelPlatesResonator(ParallelHorizontalPlatesResonator):
    pass

class ParallelVerticalPlatesResonator(Resonator):
    '''Broad-band resonator for vertical parallel plates.'''
    def __init__(self, R_shunt, frequency, Q, n_turns_wake=1,
                 *args, **kwargs):
        """ Special case of parallel plate resonator. """
        Yokoya_X1 = np.pi**2 / 12.
        Yokoya_Y1 = np.pi**2 / 24.
        Yokoya_X2 = np.pi**2 / 24.
        Yokoya_Y2 = -np.pi**2 / 24.
        switch_Z  = False

        super(ParallelVerticalPlatesResonator, self).__init__(
            R_shunt, frequency, Q, Yokoya_X1, Yokoya_Y1,
            Yokoya_X2, Yokoya_Y2, switch_Z, n_turns_wake, *args, **kwargs)

class ResistiveWall(WakeSource):
    """ Class to describe the wake functions originating from a
    resistive wall impedance. """

    def __init__(self, pipe_radius, resistive_wall_length, conductivity,
                 dt_min, Yokoya_X1, Yokoya_Y1, Yokoya_X2, Yokoya_Y2,
                 n_turns_wake=1, *args, **kwargs):
        """ General constructor to create a ResistiveWall WakeSource
        object describing the wake functions of a resistive wall
        impedance.
        The parameter 'n_turns_wake' defines how many turns are
        considered for the multiturn wakes. It is 1 by default, i.e.
        multiturn wakes are off. """
        super(ResistiveWall, self).__init__(*args, **kwargs)

        self.pipe_radius = np.array([pipe_radius]).flatten()
        self.resistive_wall_length = resistive_wall_length
        self.conductivity = conductivity
        self.dt_min = dt_min

        self.Yokoya_X1 = Yokoya_X1
        self.Yokoya_Y1 = Yokoya_Y1
        self.Yokoya_X2 = Yokoya_X2
        self.Yokoya_Y2 = Yokoya_Y2
        self.n_turns_wake = n_turns_wake

    def get_wake_kicks(self, slicer):
        """ Factory method. Creates instances of the appropriate
        WakeKick objects for the ResistiveWall WakeSource with the
        specified parameters. A WakeKick object is instantiated only if
        the corresponding Yokoya factor is non-zero. The WakeKick
        objects are returned as a list wake_kicks. """
        wake_kicks = []

        # Dipole wake kick x.
        if self.Yokoya_X1:
            wake_function = self.function_transverse(self.Yokoya_X1)
            wake_kicks.append(DipoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        # Quadrupole wake kick x.
        if self.Yokoya_X2:
            wake_function = self.function_transverse(self.Yokoya_X2)
            wake_kicks.append(QuadrupoleWakeKickX(
                wake_function, slicer, self.n_turns_wake))

        # Dipole wake kick y.
        if self.Yokoya_Y1:
            wake_function = self.function_transverse(self.Yokoya_Y1)
            wake_kicks.append(DipoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        # Quadrupole wake kick y.
        if self.Yokoya_Y2:
            wake_function = self.function_transverse(self.Yokoya_Y2)
            wake_kicks.append(QuadrupoleWakeKickY(
                wake_function, slicer, self.n_turns_wake))

        return wake_kicks

    def function_transverse(self, Yokoya_factor):
        """ Define the wake function (transverse) of a resistive wall
        with the given parameters. """
        mu_r = 1
        # The impedance of free space [Ohm]
        Z_0 = 119.9169832 * np.pi

        def wake(dt, *args, **kwargs):
            y = (Yokoya_factor * (np.sign(dt + np.abs(self.dt_min)) - 1) / 2. *
                 np.sqrt(kwargs['beta']) * self.resistive_wall_length / np.pi /
                 self.pipe_radius**3 * np.sqrt(-mu_r / np.pi /
                 self.conductivity / dt.clip(max=-abs(self.dt_min))))*np.sqrt(Z_0*c)
            return y
        return wake

class CircularResistiveWall(ResistiveWall):
    '''Circular resistive wall.'''
    def __init__(self, pipe_radius, resistive_wall_length, conductivity,
                 dt_min, n_turns_wake=1, *args, **kwargs):
        """ Special case of a circular resistive wall. """
        Yokoya_X1 = 1.
        Yokoya_Y1 = 1.
        Yokoya_X2 = 0.
        Yokoya_Y2 = 0.

        super(CircularResistiveWall, self).__init__(
            pipe_radius, resistive_wall_length, conductivity, dt_min,
            Yokoya_X1, Yokoya_Y1, Yokoya_X2, Yokoya_Y2, n_turns_wake,
            *args, **kwargs)

class ParallelHorizontalPlatesResistiveWall(ResistiveWall):
    '''Resistive wall impedance for horizontal parallel plates.'''
    def __init__(self, pipe_radius, resistive_wall_length, conductivity,
                 dt_min, n_turns_wake=1, *args, **kwargs):
        """ Special case of a parallel plates resistive wall. """
        Yokoya_X1 = np.pi**2 / 24.
        Yokoya_Y1 = np.pi**2 / 12.
        Yokoya_X2 = -np.pi**2 / 24.
        Yokoya_Y2 = np.pi**2 / 24.

        super(ParallelHorizontalPlatesResistiveWall, self).__init__(
            pipe_radius, resistive_wall_length, conductivity, dt_min,
            Yokoya_X1, Yokoya_Y1, Yokoya_X2, Yokoya_Y2, n_turns_wake,
            *args, **kwargs)

@deprecated('--> "ParallelPlatesResistiveWall" will be removed '
            'in the near future. '
            'Use "ParallelHorizontalPlatesResistiveWall" instead.\n')
class ParallelPlatesResistiveWall(ParallelHorizontalPlatesResistiveWall):
    pass

class ParallelVerticalPlatesResistiveWall(Resonator):
    '''Resistive wall impedance for vertical parallel plates.'''
    def __init__(self, pipe_radius, resistive_wall_length, conductivity,
                 dt_min, n_turns_wake=1, *args, **kwargs):
        """ Special case of a parallel plates resistive wall. """
        Yokoya_X1 = np.pi**2 / 12.
        Yokoya_Y1 = np.pi**2 / 24.
        Yokoya_X2 = np.pi**2 / 24.
        Yokoya_Y2 = -np.pi**2 / 24.

        super(ParallelVerticalPlatesResistiveWall, self).__init__(
            pipe_radius, resistive_wall_length, conductivity, dt_min,
            Yokoya_X1, Yokoya_Y1, Yokoya_X2, Yokoya_Y2, n_turns_wake,
            *args, **kwargs)