content: from qutip import Qobj, qeye_like
  from .cy.dysolve import cy_compute_integrals
  from numpy.typing import ArrayLike
  import numpy as np
  import scipy as sp
  from numbers import Number
  import itertools
  
  
  __all__ = ['DysolvePropagator', 'dysolve_propagator']
  
  
  class DysolvePropagator:
      """
      A generator of propagator using Dysolve.
      https://arxiv.org/abs/2012.09282
  
      Parameters
      ----------
      H_0 : Qobj
          The base hamiltonian of the system.
  
      X : Qobj
          A cosine perturbation applied on the system.
  
      omega : float
          The frequency of the cosine perturbation.
  
      options : dict, optional
          Extra parameters.
  
          - "max_order"
  
              A given integer to indicate the highest order of
              approximation used to compute the propagators (default is 4).
              This corresponds to n in eq. (4) of Ref.
  
          - "a_tol"
  
              The absolute tolerance used when computing the propagators
              (default is 1e-10).
  
          - "max_dt"
  
              The maximum time increment used when computing propagators
              (default is 0.1).
  
      Notes
      -----
      The system's hamiltonian must be of the form
      H = H_0 + cos(omega*t)X for Dysolve to work.
  
      For the moment, only a cosine perturbation is allowed. Dysolve can
      manage more exotic perturbations, but this is not implemented yet.
  
      .. note:: Experimental.
  
      """
  
      def __init__(
          self,
          H_0: Qobj,
          X: Qobj,
          omega: float,
          options: dict[str] = None,
      ):
          # System
          self._eigenenergies, self._basis = H_0.eigenstates()
          self._H_0 = H_0.transform(self._basis)
          self._X = X.transform(self._basis)
          self._elems = self._X.full().flatten()
          self._omega = omega
  
          # Options
          if options is None:
              self.max_order = 4
              self.a_tol = 1e-10
              self.max_dt = 0.1
          else:
              self.max_order = options.get('max_order', 4)
              self.max_dt = options.get('max_dt', 0.1)
              self.a_tol = options.get('a_tol', 1e-10)
  
          # Memoization
          self._dt_Sns = {}
  
          # Time propagator
          self.U = None
  
      def __call__(self, t_f: float, t_i: float = 0.0) -> Qobj:
          """
          Computes the propagator from t_i to t_f. If t_i is not provided,
          computes the propagator from 0 to t_f.
  
          Parameters
          ----------
          t_f : float
              Final time of the evolution.
  
          t_i : float, default = 0.0
              Initial time of the evolution.
  
          Returns
          -------
          U : Qobj
              The propagator U(t_f, t_i) from t_i to t_f.
  
          Notes
          -----
          If t_f - t_i > max_dt, splits the evolution into smaller ones
          to then reconstruct U(t_f, t_i).
  
          Memoization is used. First call may be slow but the next calls
          should be faster.
  
          """
          time_diff = t_f - t_i
          dt = self.max_dt * np.sign(time_diff)
          n_steps = abs(int(time_diff / self.max_dt))
  
          U = self._compute_subprop(t_i, dt)
  
          for j in range(1, n_steps):
              U = self._compute_subprop(t_i + j*dt, dt) @ U
  
          remaining = time_diff - n_steps*dt
          if abs(remaining) > self.a_tol:
              dt = remaining
              U = self._compute_subprop(t_f - dt, dt) @ U
  
          self.U = Qobj(U, self._H_0._dims, copy=False).transform(
              self._basis, True
          )
  
          return self.U
  
      def _update_matrix_elements(self, current: ArrayLike) -> ArrayLike:
          """
          Reuses the current matrix elements (order n-1) to compute the
          matrix elements for the order n.
  
          Parameters
          ----------
          current : ArrayLike
              The current matrix elements (for the order n-1)..
  
          Returns
          -------
          matrix_elements : ArrayLike
              The new matrix elements for the order n.
          """
          if current is None:
              return self._elems
          else:
              shape = self._X.shape[0]
              a = np.tile(current, shape)
              b = np.repeat(self._elems, len(current)//shape)
              return a * b
  
          # WIP:
          # The following is an attempt to use scipy.sparse to store "current".
          # There can be a lot of zeros, so scipy.sparse could be useful here.
          # This involves more operations than the implementation above because
          # "current" is a vector and scipy.sparse only accepts matrices.
          # A better approach is necessary to use the full potential of
          # scipy.sparse.
  
          # if current is None:
          #     return elems
          # else:
          #     x_shape = self._X.shape[0]
          #     a = sp.sparse.hstack([current]*x_shape)
          #     b = sp.sparse.vstack(
          #             [elems]*(current.shape[1]//x_shape)
          #     ).transpose().reshape(a.shape)
          #     return a.multiply(b)
  
      def _compute_Sns(self, dt: float) -> dict:
          """
          Computes Sns for each omega vector. This implements a similar equation
          to eq. (14) in Ref, but the function "f" is not used to avoid dealing
          explicitly with limits.
  
          Parameters
          ----------
          dt : float
              The time increment.
  
          Returns
          -------
          Sns : dict
              Sns for each omega vector. key = order with the result for each
              omega vector.
  
          """
          if dt in self._dt_Sns:
              return self._dt_Sns[dt]
  
          else:
              Sns = {}
              length = len(self._eigenenergies)
              exp_H_0 = (-1j*dt*self._H_0).expm().full()
              current_matrix_elements = None
  
              Sns[0] = exp_H_0
  
              for n in range(1, self.max_order + 1):
                  omega_vectors = np.fromiter(
                      itertools.product([self._omega, -self._omega], repeat=n),
                      np.dtype((float, (n,)))
                  )
  
                  lambdas = np.fromiter(
                      itertools.product(self._eigenenergies, repeat=n + 1),
                      np.dtype((float, (n+1,)))
                  )
                  diff_lambdas = -np.diff(lambdas)[:, ::-1]
  
                  ket_bra_idx = np.vstack(
                      (np.repeat(np.arange(0, length), length**n),
                       np.tile(np.arange(0, length), length**n))
                  ).T
  
                  Sn = np.zeros((len(omega_vectors), length, length),
                                dtype=np.complex128
                                )
  
                  # Compute matrix elements
                  current_matrix_elements = self._update_matrix_elements(
                      current_matrix_elements
                  )
  
                  for i, omega_vector in enumerate(omega_vectors):
                      # Compute integrals
                      ls_ws = omega_vector + diff_lambdas
  
                      for j, ws in enumerate(ls_ws):
                          if current_matrix_elements[j] != 0:
                              x = cy_compute_integrals(
                                  ws, dt) * current_matrix_elements[j]
                              Sn[i, ket_bra_idx[j, 0], ket_bra_idx[j, 1]] += x
  
                      Sn[i] *= (-1j / 2) ** n
                      Sn[i] = exp_H_0 @ Sn[i]
  
                  Sns[n] = Sn
  
              self._dt_Sns[dt] = Sns
              return Sns
  
      def _compute_subprop(self, current_time: float, dt: float) -> ArrayLike:
          """
          Computes a subpropagator U(current_time + dt, current_time).
  
          Parameters
          ----------
          current_time : float
              The starting time of the evolution. Can be positive or negative.
  
          dt : float
              The time increment.
  
          Returns
          -------
          subpropagator : ArrayLike
              U(current_time + dt, current_time).
  
          """
          Sns = self._compute_Sns(dt)
  
          subpropagator = np.zeros(
              (len(self._eigenenergies), len(self._eigenenergies)),
              dtype=np.complex128
          )
          subpropagator += Sns[0]
  
          for n in range(1, self.max_order + 1):
              omega_vectors = np.fromiter(
                  itertools.product([self._omega, -self._omega], repeat=n),
                  np.dtype((float, (n,)))
              )
              subpropagator += sum(
                  Sns[n] * np.exp(1j*np.sum(omega_vectors, axis=1)*current_time)
                  [:, np.newaxis, np.newaxis]
              )
  
          return subpropagator
  
  
  def dysolve_propagator(
          H_0: Qobj,
          X: Qobj,
          omega: float,
          t: float | list[float],
          options: dict[str] = None
  ) -> Qobj | list[Qobj]:
      """
      A generator of propagator(s) using Dysolve.
      https://arxiv.org/abs/2012.09282.
  
      Parameters
      ----------
      H_0 : Qobj
          The hamiltonian of the system.
  
      X : Qobj
          A cosine perturbation applied on the system.
  
      omega : float
          The frequency of the cosine perturbation.
  
      t : float | list[float]
          Time or list of times for which to evaluate the propagator(s). If t
          is a single number, the propagator from 0 to t is computed. When
          t is a list, the propagators from the first time to each elements in
          t is returned. In that case, the first output will always be the
          identity matrix. Also, in that case, have the same time increment in
          between elements for better performance.
  
      options : dict, optional
          Extra parameters.
  
          - "max_order"
  
              A given integer to indicate the highest order of
              approximation used to compute the propagators (default is 4).
              This corresponds to n in eq. (4) of Ref.
  
          - "a_tol"
  
              The absolute tolerance used when computing the propagators
              (default is 1e-10).
  
          - "max_dt"
  
              The maximum time increment used when computing propagators
              (default is 0.1).
  
      Returns
      -------
      Us : Qobj | list[Qobj]
          The time evolution propagator U(t,0) if t is a single number or else
          a list of propagators [U(t[i], t[0])] for all elements t[i] in t.
  
      Notes
      -----
      The system's hamiltonian must be of the form
      H = H_0 + cos(omega*t)X for Dysolve to work.
  
      For the moment, only a cosine perturbation is allowed. Dysolve can
      manage more exotic perturbations, but this is not implemented yet.
  
      .. note:: Experimental.
  
      """
      if isinstance(t, Number):
          dysolve = DysolvePropagator(H_0, X, omega, options)
          return dysolve(t)
  
      else:
          Us = []
          Us.append(qeye_like(H_0))  # U(t_0, t_0) = identity
  
          dysolve = DysolvePropagator(H_0, X, omega, options)
          for i in range(len(t[:-1])):  # Compute individual U(t[i+1], t[i])
              U = dysolve(t[i+1], t[i])
              Us.append(U)
  
          for i in range(1, len(Us)):  # [U(t[i], t[0])]
              Us[i] = Us[i] @ Us[i - 1]
  
      return Us
  