content: # -*- coding: utf-8 -*-
  # @author: Arne L. Grimsmo
  # @email1: arne.grimsmo@gmail.com
  # @organization: University of Sherbrooke
  
  """
  This module contains an implementation of the non-Markovian transfer tensor
  method (TTM), introduced in [1].
  
  [1] Javier Cerrillo and Jianshu Cao, Phys. Rev. Lett 112, 110401 (2014)
  """
  
  import numpy as np
  import time
  from qutip import vector_to_operator, operator_to_vector, Result
  
  
  def ttmsolve(dynmaps, state0, times, e_ops=(), num_learning=0, options=None):
      """
      Expand time-evolution using the Transfer Tensor Method [1]_, based on a set
      of precomputed dynamical maps.
  
      Parameters
      ----------
      dynmaps : list of :class:`.Qobj`, callable
          List of precomputed dynamical maps (superoperators) for the first times
          of ``times`` or a callback function that returns the superoperator at a
          given time.
  
      state0 : :class:`.Qobj`
          Initial density matrix or state vector (ket).
  
      times : array_like
          List of times :math:`t_n` at which to compute results.
          Must be uniformily spaced.
  
      e_ops : :class:`.Qobj`, callable, or list, optional
          Single operator or list of operators for which to evaluate
          expectation values or callable or list of callable.
          Callable signature must be, `f(t: float, state: Qobj)`.
          See :func:`expect` for more detail of operator expectation.
  
      num_learning : int, default: 0
          Number of times used to construct the dynmaps operators when
          ``dynmaps`` is a callable.
  
      options : dictionary, optional
          Dictionary of options for the solver.
  
          - store_final_state : bool
            Whether or not to store the final state of the evolution in the
            result class.
          - store_states : bool, None
            Whether or not to store the state vectors or density matrices.
            On `None` the states will be saved if no expectation operators are
            given.
          - normalize_output : bool
            Normalize output state to hide ODE numerical errors.
          - threshold : float
            Threshold for halting. Halts if  :math:`||T_{n}-T_{n-1}||` is below
            treshold.
  
      Returns
      -------
      output: :class:`.Result`
          An instance of the class :class:`.Result`.
  
      .. [1] Javier Cerrillo and Jianshu Cao, Phys. Rev. Lett 112, 110401 (2014)
      """
  
      opt = {
          "store_final_state": False,
          "store_states": None,
          "normalize_output": True,
          "threshold": 0.0,
          "num_learning": 0,
      }
      if options:
          opt.update(options)
  
      if not np.allclose(np.diff(times), times[1] - times[0]):
          raise ValueError("The time should be uniformily distributed.")
  
      if callable(dynmaps):
          if num_learning <= 0:
              raise ValueError(
                  "When dynmaps is a callable, options['num_learning'] must be "
                  "the number of dynamical maps to compute."
              )
          dynmaps = [dynmaps(t) for t in times[:num_learning]]
  
      if (
          not dynmaps
          or not dynmaps[0].issuper
          or not all(dmap._dims == dynmaps[0]._dims for dmap in dynmaps)
      ):
          raise ValueError("`dynmaps` entries must be super operators.")
  
      start = time.time()
      tensors, diff = _generatetensors(dynmaps, opt["threshold"])
      end = time.time()
  
      stats = {
          "preparation time": end - start,
          "run time": 0.0,
          "ttmconvergence": diff,
          "num_tensor": len(tensors),
      }
      if state0.isket:
          state0 = state0.proj()
  
      if state0.isoper:
          # vectorize density matrix
          rho0vec = operator_to_vector(state0)
          restore = vector_to_operator
      else:
          # state0 might be a super in which case we should not vectorize
          rho0vec = state0
          restore = lambda state: state
  
      K = len(tensors)
      start = time.time()
      results = Result(e_ops, opt, solver="ttmsolve", stats=stats)
      states = [rho0vec]
      results.add(times[0], state0)
  
      for n in range(1, len(times)):
          # Set current state
          state = 0
          for j in range(1, min(K, n + 1)):
              tmp = tensors[j] @ states[n - j]
              state = tmp + state
          # Append state to all states
          states.append(state)
          results.add(times[n], restore(state))
      end = time.time()
      stats["run time"] = end - start
  
      return results
  
  
  def _generatetensors(dynmaps, threshold):
      r"""
      Generate the tensors :math:`T_1,\dots,T_K` from the dynamical maps
      :math:`E(t_k)`.
  
      A stationary process is assumed, i.e., :math:`T_{n,k} = T_{n-k}`.
  
      Parameters
      ----------
      dynmaps : list of :class:`.Qobj`
          List of precomputed dynamical maps (superoperators) at the times
          specified in `learningtimes`.
  
      threshold : float
          Threshold for halting. Halts if  :math:`||T_{n}-T_{n-1}||` is below
          treshold.
  
      Returns
      -------
      Tensors, diffs: list of :class:`.Qobj.`
          A list of transfer tensors :math:`T_1,\dots,T_K`
      """
      Tensors = []
      diff = [0.0]
      for n in range(len(dynmaps)):
          T = dynmaps[n]
          for m in range(1, n):
              T -= Tensors[n - m] @ dynmaps[m]
          Tensors.append(T)
          if n > 1:
              diff.append((Tensors[-1] - Tensors[-2]).norm())
              if diff[-1] < threshold:
                  # Below threshold for truncation
                  break
      return Tensors, diff
  