content: # -*- 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
  
  from qutip import (
      sprepost,
      Qobj,
      spre,
      spost,
      liouvillian,
      qeye,
      mesolve,
      propagator,
      composite,
      isket,
      ket2dm,
      tensor_contract,
  )
  
  
  __all__ = ["MemoryCascade"]
  
  
  class MemoryCascade:
      """Class for running memory cascade simulations of open quantum systems
      with time-delayed coherent feedback.
  
      Attributes
      ----------
      H_S : :class:`.Qobj`
          System Hamiltonian (can also be a Liouvillian)
  
      L1 : :class:`.Qobj` / list of :class:`.Qobj`
          System operators coupling into the feedback loop. Can be a single
          operator or a list of operators.
  
      L2 : :class:`.Qobj` / list of :class:`.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:`.Qobj` / list of :class:`.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).
  
      options : dict
          Generic solver options.
      """
  
      def __init__(
          self,
          H_S,
          L1,
          L2,
          S_matrix=None,
          c_ops_markov=None,
          integrator="propagator",
          options=None,
      ):
  
          if options is None:
              self.options = {"progress_bar": False}
          else:
              self.options = options
  
          self.H_S = H_S
          self.sysdims = H_S.dims
          if isinstance(L1, Qobj):
              self.L1 = [L1]
          else:
              self.L1 = L1
          if isinstance(L2, 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, 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 = qeye(H_S.shape[0])
          self.Id.dims = self.sysdims
          self.Id = sprepost(self.Id, self.Id)
          self.store_states = self.options.get("store_states", False)
          self.integrator = integrator
          self._generators = {}
  
      def generator(self, k):
          if k not in self._generators:
              self._generators[k] = _generator(
                  k, self.H_S, self.L1, self.L2, self.S_matrix, self.c_ops_markov
              )
          return self._generators[k]
  
      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:`.Qobj`
              time-propagator for reduced system dynamics
          """
          k = int(t / tau) + 1
          s = t - (k - 1) * tau
          G1 = self.generator(k)
          E0 = qeye(G1.dims[0])
          E = _integrate(
              G1, E0, 0.0, s, integrator=self.integrator, opt=self.options
          )
          if k > 1:
              G2 = self.generator(k - 1)
              G2 = composite(G2, self.Id)
              E = _integrate(
                  G2, E, s, tau, integrator=self.integrator, opt=self.options
              )
  
          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:`.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:`.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:`.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 = self.generator(kmax)
          sprev = 0.0
          E0 = qeye(G1.dims[0])
          E = E0
          for i, s in enumerate(slist):
              E = _integrate(
                  G1, E, sprev, s, integrator=self.integrator, opt=self.options
              )
              l2 = _localop(c2, klist[i], kmax)
              if klist[i] == 1:
                  l1 = l2 * 0.0
              else:
                  l1 = _localop(c1, klist[i] - 1, kmax)
              if blist[i] == 0:
                  superop = self.Id
              elif blist[i] == 1:
                  superop = spre(l1 + l2)
              elif blist[i] == 2:
                  superop = spost(l1.dag() + l2.dag())
              elif blist[i] == 3:
                  superop = spre(l1)
              elif blist[i] == 4:
                  superop = spost(l1.dag())
              else:
                  raise ValueError(
                      "Allowed values in blist are 0, 1, 2, 3 and 4."
                  )
  
              E = superop @ E
              sprev = s
          E = _integrate(
              G1, E, slist[-1], tau, integrator=self.integrator, opt=self.options
          )
  
          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:`.Qobj`
              initial density matrix or state vector (ket)
  
          t : float
              current time
  
          tau : float
              time-delay
  
          Returns
          -------
          : :class:`.Qobj`
              density matrix at time :math:`t`
          """
          if isket(rho0):
              rho0 = ket2dm(rho0)
  
          E = self.propagator(t, tau)
          return E(rho0)
  
      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:`.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:`.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:`.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
          """
          if isket(rho0):
              rho0 = ket2dm(rho0)
  
          E = self.outfieldpropagator(blist, tlist, tau)
          return (E(rho0)).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")
      out = op
      if l > 1:
          out = qeye(op.dims[0] * (l - 1)) & out
      if l < k:
          out = out & qeye(op.dims[0] * (k - l))
  
      return out
  
  
  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 = 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 = qeye(H.dims[0][0])
      Id = sprepost(id, id)
      if S is None:
          S = np.identity(len(L1))
      # create Lindbladian
  
      # first system
      L = liouvillian(None, [_localop(c, 1, k) for c in L2])
      for l in range(1, k):
          # Bare Hamiltonian
          Hl = _localop(H, l, k)
          L += liouvillian(Hl, [])
          # Markovian Decay channels
          if c_ops_markov is not None:
              for c in c_ops_markov:
                  cl = _localop(c, l, k)
                  L += 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 += liouvillian(Hcasc, [c for c in Lvec])
      # last system
      L += 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 += liouvillian(None, [cl])
      # return generator
      return L
  
  
  def _integrate(L, E0, ti, tf, integrator="propagator", opt=None):
      """
      Basic ode integrator
      """
      opt = opt or {}
      if tf > ti:
          if integrator == "mesolve":
              opt["store_final_state"] = True
              sol = mesolve(L, E0, [ti, tf], options=opt)
              return sol.final_state
          elif integrator == "propagator":
              return propagator(L, (tf - ti), options=opt) @ E0
          else:
              raise ValueError(
                  'integrator keyword must be either "propagator" or "mesolve"'
              )
      else:
          return E0
  