# Source code for qutip.solver.propagator

    # Required for Sphinx to follow autodoc_type_aliases
    from __future__ import annotations

    __all__ = ['Propagator', 'propagator', 'propagator_steadystate']

    import numbers
    import numpy as np

    from .. import Qobj, qeye, qeye_like, unstack_columns, QobjEvo, liouvillian
    from ..core import data as _data
    from ..typing import QobjEvoLike
    from .mesolve import mesolve, MESolver
    from .sesolve import sesolve, SESolver
    from .heom.bofin_solvers import HEOMSolver
    from .solver_base import Solver
    from .multitraj import MultiTrajSolver
    from numbers import Number
    from typing import Any
    from qutip.settings import settings

    def propagator_piecewise(
        H: QobjEvo,
        tlist: list[Number],
        piecewise_t: list[Number],
        c_ops: list[QobjEvo] = None,
        args: dict[str, Any] = None,
    ) -> Qobj | list[Qobj]:
        r"""
        Calculate propagators for piecewise constant Hamiltonians or
        Liouvillians.

        This function efficiently computes propagators when the Hamiltonian or
        Liouvillian is piecewise constant, i.e., it remains constant over
        certain time intervals and changes discontinuously at specified times.
        By exponentiating the generator directly on each constant interval,
        this approach is significantly faster than numerically integrating the
        differential equation.

        Parameters
        ----------
        H : :obj:`.QobjEvo`
            Time-dependent Hamiltonian or Liouvillian.

        tlist : list of float
            List of times at which to evaluate the propagator.

        piecewise_t : list of float
            Times where the Hamiltonian or Liouvillian or collapse operators change
            discontinuously.  These times define the boundaries between constant
            intervals.
            Times outside the range ``(tlist[0], tlist[-1]]`` are ignored.

        c_ops : list of :obj:`.QobjEvo`, optional
            List of collapse operators for open quantum systems.
            Can be piecewise constant.

        args : dict, optional
            Dictionary of parameters to pass to the Hamiltonian and collapse
            operator functions.

        Returns
        -------
        propagators : list of :obj:`.Qobj`
            List of propagators :math:`U(t)` evaluated at each time in ``tlist``,
            such that :math:`\psi(t) = U(t)\psi(t_0)` for state vectors or
            :math:`\rho_{\mathrm{vec}}(t) = U(t)\rho_{\mathrm{vec}}(t_0)` for
            vectorized density matrices.
        """

        piecewise_times = {tt for tt in piecewise_t if tlist[0] < tt <= tlist[-1]}
        eval_times = set(tlist)
        times = sorted(eval_times | piecewise_times)
        atol = settings.core["atol"]

        out = []
        prev = times[0]
        prev_dt = None
        dU = None

        for nxt in times[1:]:
            dt = nxt - prev

            # Evaluate at midpoint to avoid discontinuities at boundaries
            t_eval = (prev + nxt) / 2
            H_step = H(t_eval, args)
            if c_ops:
                c_ops_q = [op(t_eval, args) for op in c_ops]
                gen = liouvillian(H_step, c_ops_q)
            else:
                gen = H_step if H_step.issuper else -1j * H_step

            if prev == times[0]:
                U = qeye_like(gen)
                if prev in eval_times:
                    out.append(U)

            cannot_reuse = (
                dU is None
                or prev_dt is None
                or abs(dt - prev_dt) > atol
                or prev in piecewise_times  # Moved past a switching point
            )

            if cannot_reuse:
                dU = (gen * dt).expm()
            U = dU @ U
            prev_dt = dt

            if nxt in eval_times:
                out.append(U)
            prev = nxt
        return out

    def propagator(
        H: QobjEvoLike,
        t: Number | list[Number],
        c_ops: QobjEvoLike | list[QobjEvoLike] = None,
        args: dict[str, Any] = None,
        options: dict[str, Any] = None,
        piecewise_t: list[Number] | None = None,
        **kwargs,
    ) -> Qobj | list[Qobj]:
        r"""
        Calculate the propagator U(t) for the density matrix or wave function such
        that :math:`\psi(t) = U(t)\psi(0)` or
        :math:`\rho_{\mathrm vec}(t) = U(t) \rho_{\mathrm vec}(0)`
        where :math:`\rho_{\mathrm vec}` is the vector representation of the
        density matrix.

        Parameters
        ----------
        H : :obj:`.Qobj`, :obj:`.QobjEvo`, :obj:`.QobjEvo` compatible format
            Possibly time-dependent system Liouvillian or Hamiltonian as a Qobj or
            QobjEvo. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`] or callable
            that can be made into :obj:`.QobjEvo` are also accepted.

        t : float or array-like
            Time or list of times for which to evaluate the propagator. If a single
            time ``t`` is passed, the propagator from ``0`` to ``t`` is computed.
            When ``t`` is a list, the propagators from the first time in the list
            to each elements in ``t`` is returned. In that case, the first output
            will always be the identity matrix.

        c_ops : list, optional
            List of collapse operators as Qobj, QobjEvo or list that can be made
            into QobjEvo.

        args : dictionary, optional
            Parameters to callback functions for time-dependent Hamiltonians and
            collapse operators.

        options : dict, optional
            Options for the solver.

        piecewise_t : list, optional
            Times where the Hamiltonian or Liouvillian change values when
            they are piecewise constant. Providing these allows for a faster
            computation by exponentiating the Liouvillian or Hamiltonian directly
            on each interval.

        **kwargs :
            Extra parameters to use when creating the
            :obj:`.QobjEvo` from a list format ``H``. The most common are ``tlist``
            and ``order`` for array-based time dependance. See :obj:`.QobjEvo` for
            the full list.

        Returns
        -------
        U : :obj:`.Qobj`, list
            Instance representing the propagator(s) :math:`U(t)`. Return a single
            Qobj when ``t`` is a number or a list when ``t`` is a list.

        Notes
        -----
        Unlike :func:`.sesolve` or :func:`.mesolve`, the output times in ``t`` are
        not used for array time dependent system. ``tlist`` must be passed as a
        keyword argument in those case. ``tlist`` and ``t`` can have different
        length and values.

        """
        if isinstance(t, numbers.Real):
            tlist = [0, t]
            list_output = False
        else:
            tlist = t
            list_output = True

        if not isinstance(H, (Qobj, QobjEvo)):
            H = QobjEvo(H, args=args, **kwargs)

        if isinstance(c_ops, list):
            c_ops = [QobjEvo(op, args=args, **kwargs) for op in c_ops]
        elif c_ops is not None:
            c_ops = [QobjEvo(c_ops, args=args, **kwargs)]

        if piecewise_t is not None:
            out = propagator_piecewise(H, tlist, piecewise_t, c_ops, args)
            return out if list_output else out[-1]

        if c_ops:
            H = liouvillian(H, c_ops)

        U0 = qeye_like(H)

        if H.issuper:
            out = mesolve(H, U0, tlist, args=args, options=options).states
        else:
            out = sesolve(H, U0, tlist, args=args, options=options).states

        if list_output:
            return out
        else:
            return out[-1]

    def propagator_steadystate(U: Qobj) -> Qobj:
        r"""Find the steady state for successive applications of the propagator
        :math:`U`.

        Parameters
        ----------
        U : :obj:`.Qobj`
            Operator representing the propagator.

        Returns
        -------
        a : :obj:`.Qobj`
            Instance representing the steady-state density matrix.
        """
        evals, estates = U.eigenstates()
        shifted_vals = np.abs(evals - 1.0)
        ev_idx = np.argmin(shifted_vals)
        rho_data = unstack_columns(estates[ev_idx].data)
        rho_data = _data.mul(rho_data, 0.5 / _data.trace(rho_data))
        return Qobj(_data.add(rho_data, _data.adjoint(rho_data)),
                    dims=U._dims[0].oper,
                    isherm=True,
                    copy=False)

    class Propagator:
        """
        A generator of propagator for a system.

        Usage:

            U = Propagator(H, c_ops)

            psi_t = U(t) @ psi_0

        Save some previously computed propagator are stored to speed up subsequent
        computation. Changing ``args`` will erase these stored probagator.

        Parameters
        ----------
        system : :obj:`.Qobj`, :obj:`.QobjEvo`, :class:`.Solver`
            Possibly time-dependent system driving the evolution, either already
            packaged in a solver, such as :class:`.SESolver` or :class:`.BRSolver`,
            or the Liouvillian or Hamiltonian as a :obj:`.Qobj`,
            :obj:`.QobjEvo`. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`]
            or callable that can be made into :obj:`.QobjEvo` are also accepted.

            Solvers that run non-deterministacilly, such as :class:`.MCSolver`, are
            not supported.

        c_ops : list, optional
            List of :obj:`.Qobj` or :obj:`.QobjEvo` collapse operators.

        args : dictionary, optional
            Parameters to callback functions for time-dependent Hamiltonians and
            collapse operators.

        options : dict, optional
            Options for the solver.

        memoize : int, default: 10
            Max number of propagator to save.

        tol : float, default: 1e-14
            Absolute tolerance for the time. If a previous propagator was computed
            at a time within tolerance, that propagator will be returned.

        Notes
        -----
        The :class:`Propagator` is not a :obj:`.QobjEvo` so
        it cannot be used for operations with :obj:`.Qobj` or
        :obj:`.QobjEvo`. It can be made into a
        :obj:`.QobjEvo` with ::

            U = QobjEvo(Propagator(H))

        """
        def __init__(
            self,
            system: Qobj | QobjEvo | Solver,
            *,
            c_ops: QobjEvoLike | list[QobjEvoLike] = None,
            args: dict[str, Any] = None,
            options: dict[str, Any] = None,
            memoize: int = 10,
            tol: float = 1e-14,
        ):
            if isinstance(system, MultiTrajSolver):
                raise TypeError("Non-deterministic solvers cannot be used "
                                "as a propagator system")
            elif isinstance(system, HEOMSolver):
                raise NotImplementedError(
                    "HEOM is not supported by Propagator. "
                    "Please, tell us on GitHub issues if you need it!"
                )
            elif isinstance(system, Solver):
                self.solver = system
            else:
                Hevo = QobjEvo(system, args=args)
                c_ops = c_ops if c_ops is not None else []
                c_ops = [QobjEvo(op, args=args) for op in c_ops]
                if Hevo.issuper or c_ops:
                    self.solver = MESolver(Hevo, c_ops=c_ops, options=options)
                else:
                    self.solver = SESolver(Hevo, options=options)

            self.times = [0]
            self.invs = [None]
            self.props = [qeye(self.solver._sys_dims)]
            self.solver.start(self.props[0], self.times[0])
            self.cte = self.solver.rhs.isconstant
            H_0 = self.solver.rhs(0)
            self.unitary = not H_0.issuper and H_0.isherm
            self.args = args
            self.memoize = max(3, int(memoize))
            self.tol = tol

        def _lookup_or_compute(self, t):
            """
            Get U(t) from cache or compute it.
            """
            idx = np.searchsorted(self.times, t)
            if idx < len(self.times) and abs(t-self.times[idx]) <= self.tol:
                U = self.props[idx]
            elif idx > 0 and abs(t-self.times[idx-1]) <= self.tol:
                U = self.props[idx-1]
            else:
                U = self._compute(t, idx)
                self._insert(t, U, idx)
            return U

        def __call__(self, t: float, t_start: float = 0, **args):
            """
            Get the propagator from ``t_start`` to ``t``.

            Parameters
            ----------
            t : float
                Time at which to compute the propagator.
            t_start: float [0]
                Time at which the propagator start such that:
                    ``psi[t] = U.prop(t, t_start) @ psi[t_start]``
            args : dict
                Argument to pass to a time dependent Hamiltonian.
                Updating ``args`` take effect since ``t=0`` and the new ``args``
                will be used in future call.
            """
            # We could improve it when the system is constant using U(2t) = U(t)**2
            if not self.cte and args and args != self.args:
                self.args = args
                self.solver._argument(args)
                self.times = [0]
                self.props = [qeye_like(self.props[0])]
                self.solver.start(self.props[0], self.times[0])

            if t_start:
                if t == t_start:
                    U = self._lookup_or_compute(0)
                if self.cte:
                    U = self._lookup_or_compute(t - t_start)
                else:
                    Uinv = self._inv(self._lookup_or_compute(t_start))
                    U = self._lookup_or_compute(t) @ Uinv
            else:
                U = self._lookup_or_compute(t)
            return U

        def inv(self, t: float, **args):
            """
            Get the inverse of the propagator at ``t``, such that
                ``psi_0 = U.inv(t) @ psi_t``

            Parameters
            ----------
            t : float
                Time at which to compute the propagator.
            args : dict
                Argument to pass to a time dependent Hamiltonian.
                Updating ``args`` take effect since ``t=0`` and the new ``args``
                will be used in future call.
            """
            return self._inv(self(t, **args))

        def _compute(self, t, idx):
            """
            Compute the propagator at ``t``, ``idx`` point to a pair of
            (time, propagator) close to the desired time.
            """
            t_last = self.solver._integrator.get_state(copy=False)[0]
            if self.times[idx-1] <= t_last <= t:
                U = self.solver.step(t)
            elif idx > 0:
                self.solver.start(self.props[idx-1], self.times[idx-1])
                U = self.solver.step(t)
            else:
                # Evolving backward in time is not supported by all integrator.
                self.solver.start(qeye_like(self.props[0]), t)
                Uinv = self.solver.step(self.times[idx])
                U = self._inv(Uinv) @ self.props[idx]
                self.solver.start(qeye_like(self.props[0]), 0.)
            return U

        def _inv(self, U):
            return U.dag() if self.unitary else U.inv()

        def _insert(self, t, U, idx):
            """
            Insert a new pair of (time, propagator) to the memorized states.
            """
            while len(self.times) >= self.memoize:
                rm_idx = self.memoize // 2
                if self.times[rm_idx] < t:
                    idx -= 1
                del self.times[rm_idx]
                del self.props[rm_idx]
            self.times.insert(idx, t)
            self.props.insert(idx, U)