# -*- coding: utf-8 -*-

# @author: Arne L. Grimsmo
# @email1: arne.grimsmo@gmail.com
# @organization: University of Sherbrooke

"""
This module is an implementation of the method introduced in [1], for
solving open quantum systems subject to coherent feedback with a single
discrete time-delay. This method is referred to as the ``memory cascade''
method in qutip.

[1] Arne L. Grimsmo, Phys. Rev. Lett 115, 060402 (2015)
"""

import numpy as np
import warnings

import qutip as qt


class MemoryCascade:
    """Class for running memory cascade simulations of open quantum systems
    with time-delayed coherent feedback.

    Attributes
    ----------
    H_S : :class:`qutip.Qobj`
        System Hamiltonian (can also be a Liouvillian)

    L1 : :class:`qutip.Qobj` / list of :class:`qutip.Qobj`
        System operators coupling into the feedback loop. Can be a single
        operator or a list of operators.

    L2 : :class:`qutip.Qobj` / list of :class:`qutip.Qobj`
        System operators coupling out of the feedback loop. Can be a single
        operator or a list of operators. L2 must have the same length as L1.

    S_matrix: *array*
        S matrix describing which operators in L1 are coupled to which
        operators in L2 by the feedback channel. Defaults to an n by n identity
        matrix where n is the number of elements in L1/L2.

    c_ops_markov : :class:`qutip.Qobj` / list of :class:`qutip.Qobj`
        Decay operators describing conventional Markovian decay channels.
        Can be a single operator or a list of operators.

    integrator : str {'propagator', 'mesolve'}
        Integrator method to use. Defaults to 'propagator' which tends to be 
        faster for long times (i.e., large Hilbert space).

    parallel : bool
        Run integrator in parallel if True. Only implemented for 'propagator'
        as the integrator method.

    options : :class:`qutip.solver.Options`
        Generic solver options.
    """

    def __init__(self, H_S, L1, L2, S_matrix=None, c_ops_markov=None,
                 integrator='propagator', parallel=False, options=None):

        if options is None:
            self.options = qt.Options()
        else:
            self.options = options

        self.H_S = H_S
        self.sysdims = H_S.dims
        if isinstance(L1, qt.Qobj):
            self.L1 = [L1]
        else:
            self.L1 = L1
        if isinstance(L2, qt.Qobj):
            self.L2 = [L2]
        else:
            self.L2 = L2
        if not len(self.L1) == len(self.L2):
            raise ValueError('L1 and L2 has to be of equal length.')
        if isinstance(c_ops_markov, qt.Qobj):
            self.c_ops_markov = [c_ops_markov]
        else:
            self.c_ops_markov = c_ops_markov

        if S_matrix is None:
            self.S_matrix = np.identity(len(self.L1))
        else:
            self.S_matrix = S_matrix
        # create system identity superoperator
        self.Id = qt.qeye(H_S.shape[0])
        self.Id.dims = self.sysdims
        self.Id = qt.sprepost(self.Id, self.Id)
        self.store_states = self.options.store_states
        self.integrator = integrator
        self.parallel = parallel

    def propagator(self, t, tau, notrace=False):
        """
        Compute propagator for time t and time-delay tau

        Parameters
        ----------
        t : *float*
            current time

        tau : *float*
            time-delay

        notrace : *bool* {False}
            If this optional is set to True, a propagator is returned for a
            cascade of k systems, where :math:`(k-1) tau < t < k tau`.
            If set to False (default), a generalized partial trace is performed
            and a propagator for a single system is returned.
        Returns
        -------
        : :class:`qutip.Qobj`
            time-propagator for reduced system dynamics
        """
        k = int(t/tau)+1
        s = t-(k-1)*tau
        G1, E0 = _generator(k, self.H_S, self.L1, self.L2, self.S_matrix,
                            self.c_ops_markov)
        E = _integrate(G1, E0, 0., s, integrator=self.integrator,
                       parallel=self.parallel, opt=self.options)
        if k > 1:
            G2, null = _generator(k-1, self.H_S, self.L1, self.L2,
                                  self.S_matrix, self.c_ops_markov)
            G2 = qt.composite(G2, self.Id)
            E = _integrate(G2, E, s, tau, integrator=self.integrator, 
                    parallel=self.parallel, opt=self.options)
        E.dims = E0.dims
        if not notrace:
            E = _genptrace(E, k)
        return E

    def outfieldpropagator(self, blist, tlist, tau, c1=None, c2=None,
                           notrace=False):
        r"""
        Compute propagator for computing output field expectation values
        <O_n(tn)...O_2(t2)O_1(t1)> for times t1,t2,... and
        O_i = I, b_out, b_out^\dagger, b_loop, b_loop^\dagger

        Parameters
        ----------
        blist : array_like
            List of integers specifying the field operators:
            0: I (nothing)
            1: b_out
            2: b_out^\dagger
            3: b_loop
            4: b_loop^\dagger

        tlist : array_like
            list of corresponding times t1,..,tn at which to evaluate the field
            operators

        tau : float
            time-delay

        c1 : :class:`qutip.Qobj`
            system collapse operator that couples to the in-loop field in
            question (only needs to be specified if self.L1 has more than one
            element)

        c2 : :class:`qutip.Qobj`
            system collapse operator that couples to the output field in
            question (only needs to be specified if self.L2 has more than one
            element)

        notrace : bool {False}
            If this optional is set to True, a propagator is returned for a
            cascade of k systems, where :math:`(k-1) tau < t < k tau`.
            If set to False (default), a generalized partial trace is performed
            and a propagator for a single system is returned.

        Returns
        -------
        : :class:`qutip.Qobj`
            time-propagator for computing field correlation function
        """
        if c1 is None and len(self.L1) == 1:
            c1 = self.L1[0]
        else:
            raise ValueError('Argument c1 has to be specified when more than' +
                             'one collapse operator couples to the feedback' +
                             'loop.')
        if c2 is None and len(self.L2) == 1:
            c2 = self.L2[0]
        else:
            raise ValueError('Argument c1 has to be specified when more than' +
                             'one collapse operator couples to the feedback' +
                             'loop.')
        klist = []
        slist = []
        for t in tlist:
            klist.append(int(t/tau)+1)
            slist.append(t-(klist[-1]-1)*tau)
        kmax = max(klist)
        zipped = sorted(zip(slist, klist, blist))
        slist = [s for (s, k, b) in zipped]
        klist = [k for (s, k, b) in zipped]
        blist = [b for (s, k, b) in zipped]

        G1, E0 = _generator(kmax, self.H_S, self.L1, self.L2, self.S_matrix,
                            self.c_ops_markov)
        sprev = 0.
        E = E0
        for i, s in enumerate(slist):
            E = _integrate(G1, E, sprev, s, integrator=self.integrator,
                    parallel=self.parallel, opt=self.options)
            if klist[i] == 1:
                l1 = 0.*qt.Qobj()
            else:
                l1 = _localop(c1, klist[i]-1, kmax)
            l2 = _localop(c2, klist[i], kmax)
            if blist[i] == 0:
                superop = self.Id
            elif blist[i] == 1:
                superop = qt.spre(l1+l2)
            elif blist[i] == 2:
                superop = qt.spost(l1.dag()+l2.dag())
            elif blist[i] == 3:
                superop = qt.spre(l1)
            elif blist[i] == 4:
                superop = qt.spost(l1.dag())
            else:
                raise ValueError('Allowed values in blist are 0, 1, 2, 3 ' +
                                 'and 4.')
            superop.dims = E.dims
            E = superop*E
            sprev = s
        E = _integrate(G1, E, slist[-1], tau, integrator=self.integrator,
                parallel=self.parallel, opt=self.options)

        E.dims = E0.dims
        if not notrace:
            E = _genptrace(E, kmax)
        return E

    def rhot(self, rho0, t, tau):
        """
        Compute the reduced system density matrix :math:`\\rho(t)`

        Parameters
        ----------
        rho0 : :class:`qutip.Qobj`
            initial density matrix or state vector (ket)

        t : float
            current time

        tau : float
            time-delay

        Returns
        -------
        : :class:`qutip.Qobj`
            density matrix at time :math:`t`
        """
        if qt.isket(rho0):
            rho0 = qt.ket2dm(rho0)

        E = self.propagator(t, tau)
        rhovec = qt.operator_to_vector(rho0)
        return qt.vector_to_operator(E*rhovec)

    def outfieldcorr(self, rho0, blist, tlist, tau, c1=None, c2=None):
        r"""
        Compute output field expectation value
        <O_n(tn)...O_2(t2)O_1(t1)> for times t1,t2,... and
        O_i = I, b_out, b_out^\dagger, b_loop, b_loop^\dagger


        Parameters
        ----------
        rho0 : :class:`qutip.Qobj`
            initial density matrix or state vector (ket).

        blist : array_like
            List of integers specifying the field operators:
            0: I (nothing)
            1: b_out
            2: b_out^\dagger
            3: b_loop
            4: b_loop^\dagger

        tlist : array_like
            list of corresponding times t1,..,tn at which to evaluate the field
            operators

        tau : float
            time-delay

        c1 : :class:`qutip.Qobj`
            system collapse operator that couples to the in-loop field in
            question (only needs to be specified if self.L1 has more than one
            element)

        c2 : :class:`qutip.Qobj`
            system collapse operator that couples to the output field in
            question (only needs to be specified if self.L2 has more than one
            element)

        Returns
        -------
        : complex
            expectation value of field correlation function
        """
        E = self.outfieldpropagator(blist, tlist, tau)
        rhovec = qt.operator_to_vector(rho0)
        return (qt.vector_to_operator(E*rhovec)).tr()


def _localop(op, l, k):
    """
    Create a local operator on the l'th system by tensoring
    with identity operators on all the other k-1 systems
    """
    if l < 1 or l > k:
        raise IndexError('index l out of range')
    h = op
    I = qt.qeye(op.shape[0])
    I.dims = op.dims
    for i in range(1, l):
        h = qt.tensor(I, h)
    for i in range(l+1, k+1):
        h = qt.tensor(h, I)
    return h


def _genptrace(E, k):
    """
    Perform a gneralized partial trace on a superoperator E, tracing out all
    subsystems but one.
    """
    for l in range(k-1):
        nsys = len(E.dims[0][0])
        E = qt.tensor_contract(E, (0, 2*nsys+1), (nsys, 3*nsys+1))
    return E


def _generator(k, H, L1, L2, S=None, c_ops_markov=None):
    """
    Create a Liouvillian for a cascaded chain of k system copies
    """
    id = qt.qeye(H.dims[0][0])
    Id = qt.sprepost(id, id)
    if S is None:
        S = np.identity(len(L1))
    # create Lindbladian
    L = qt.Qobj()
    E0 = Id
    # first system
    L += qt.liouvillian(None, [_localop(c, 1, k) for c in L2])
    for l in range(1, k):
        # Identiy superoperator
        E0 = qt.composite(E0, Id)
        # Bare Hamiltonian
        Hl = _localop(H, l, k)
        L += qt.liouvillian(Hl, [])
        # Markovian Decay channels
        if c_ops_markov is not None:
            for c in c_ops_markov:
                cl = _localop(c, l, k)
                L += qt.liouvillian(None, [cl])
        # Cascade coupling
        c1 = np.array([_localop(c, l, k) for c in L1])
        c2 = np.array([_localop(c, l+1, k) for c in L2])
        c2dag = np.array([c.dag() for c in c2])
        Hcasc = -0.5j*np.dot(c2dag, np.dot(S, c1))
        Hcasc += Hcasc.dag()
        Lvec = c2 + np.dot(S, c1)
        L += qt.liouvillian(Hcasc, [c for c in Lvec])
    # last system
    L += qt.liouvillian(_localop(H, k, k), [_localop(c, k, k) for c in L1])
    if c_ops_markov is not None:
        for c in c_ops_markov:
            cl = _localop(c, k, k)
            L += qt.liouvillian(None, [cl])
    E0.dims = L.dims
    # return generator and identity superop E0
    return L, E0


def _integrate(L, E0, ti, tf, integrator='propagator', parallel=False,
        opt=qt.Options()):
    """
    Basic ode integrator
    """
    if tf > ti:
        if integrator == 'mesolve':
            if parallel:
                warnings.warn('parallelization not implemented for "mesolve"')
            opt.store_final_state = True
            sol = qt.mesolve(L, E0, [ti, tf], [], [], options=opt)
            return sol.final_state
        elif integrator == 'propagator':
            return qt.propagator(L, (tf-ti), [], [], parallel=parallel,
                                 options=opt)*E0
        else:
            raise ValueError('integrator keyword must be either "propagator"' +
                              'or "mesolve"')
    else:
        return E0
