# Source code for qutip.solver.result

    """ Class for solve function results"""

    # Required for Sphinx to follow autodoc_type_aliases
    from __future__ import annotations

    from typing import TypedDict, Any, Callable
    from ..core.numpy_backend import np
    from numpy.typing import ArrayLike
    from ..core import Qobj, QobjEvo, expect

    __all__ = ["Result"]

    class _QobjExpectEop:
        """
        Pickable e_ops callable that calculates the expectation value for a given
        operator.

        Parameters
        ----------
        op : :obj:`.Qobj`
            The expectation value operator.
        """

        def __init__(self, op):
            self.op = op

        def __call__(self, t, state):
            return expect(self.op, state)

    class ExpectOp:
        """
        A result e_op (expectation operation).

        Parameters
        ----------
        op : object
            The original object used to define the e_op operation, e.g. a
            :~obj:`Qobj` or a function ``f(t, state)``.

        f : function
            A callable ``f(t, state)`` that will return the value of the e_op
            for the specified state and time.

        append : function
            A callable ``append(value)``, e.g. ``expect[k].append``, that will
            store the result of the e_ops function ``f(t, state)``.

        Attributes
        ----------
        op : object
            The original object used to define the e_op operation.
        """

        def __init__(self, op, f, append):
            self.op = op
            self._f = f
            self._append = append

        def __call__(self, t, state):
            """
            Return the expectation value for the given time, ``t`` and
            state, ``state``.
            """
            return self._f(t, state)

        def _store(self, t, state):
            """
            Store the result of the e_op function. Should only be called by
            :class:`~Result`.
            """
            self._append(self._f(t, state))

    class _BaseResult:
        """
        Common method for all ``Result``.
        """

        def __init__(self, options, *, solver=None, stats=None):
            self.solver = solver
            if stats is None:
                stats = {}
            self.stats = stats

            self._state_processors = []
            self._state_processors_require_copy = False

            # make sure not to store a reference to the solver
            options_copy = options.copy()
            if hasattr(options_copy, "_feedback"):
                options_copy._feedback = None
            self.options = options_copy
            # Almost all integrators already return a copy that is safe to use.
            self._integrator_return_copy = options.get("method", None) in [
                "adams", "lsoda", "bdf", "dop853", "diag",
                "euler", "platen", "explicit1.5",
                "milstein", "pred_corr", "taylor1.5",
                "milstein_imp", "taylor1.5_imp", "rouchon",
            ]

        def _e_ops_to_dict(self, e_ops):
            """Convert the supplied e_ops to a dictionary of Eop instances."""
            if e_ops is None:
                e_ops = {}
            elif isinstance(e_ops, (list, tuple)):
                e_ops = {k: e_op for k, e_op in enumerate(e_ops)}
            elif isinstance(e_ops, dict):
                pass
            else:
                e_ops = {0: e_ops}
            return e_ops

        def add_processor(self, f, requires_copy=False):
            """
            Append a processor ``f`` to the list of state processors.

            Parameters
            ----------
            f : function, ``f(t, state)``
                A function to be called each time a state is added to this
                result object. The state is the state passed to ``.add``, after
                applying the pre-processors, if any.

            requires_copy : bool, default False
                Whether this processor requires a copy of the state rather than
                a reference. A processor must never modify the supplied state, but
                if a processor stores the state it should set ``require_copy`` to
                true.
            """
            self._state_processors.append(f)
            self._state_processors_require_copy |= requires_copy

    class ResultOptions(TypedDict):
        store_states: bool | None
        store_final_state: bool

    class Result(_BaseResult):
        """
        Base class for storing solver results.

        Parameters
        ----------
        e_ops : :obj:`.Qobj`, :obj:`.QobjEvo`, function or list or dict of these
            The ``e_ops`` parameter defines the set of values to record at
            each time step ``t``. If an element is a :obj:`.Qobj` or
            :obj:`.QobjEvo` the value recorded is the expectation value of that
            operator given the state at ``t``. If the element is a function, ``f``,
            the value recorded is ``f(t, state)``.

            The values are recorded in the ``e_data`` and ``expect`` attributes of
            this result object. ``e_data`` is a dictionary and ``expect`` is a
            list, where each item contains the values of the corresponding
            ``e_op``.

        options : dict
            The options for this result class.

        solver : str or None
            The name of the solver generating these results.

        stats : dict or None
            The stats generated by the solver while producing these results. Note
            that the solver may update the stats directly while producing results.

        kw : dict
            Additional parameters specific to a result sub-class.

        Attributes
        ----------
        times : list
            A list of the times at which the expectation values and states were
            recorded.

        states : list of :obj:`.Qobj`
            The state at each time ``t`` (if the recording of the state was
            requested).

        final_state : :obj:`.Qobj`:
            The final state (if the recording of the final state was requested).

        expect : list of arrays of expectation values
            A list containing the values of each ``e_op``. The list is in
            the same order in which the ``e_ops`` were supplied and empty if
            no ``e_ops`` were given.

            Each element is itself a list and contains the values of the
            corresponding ``e_op``, with one value for each time in ``.times``.

            The same lists of values may be accessed via the ``.e_data`` dictionary
            and the original ``e_ops`` are available via the ``.e_ops`` attribute.

        e_data : dict
            A dictionary containing the values of each ``e_op``. If the ``e_ops``
            were supplied as a dictionary, the keys are the same as in
            that dictionary. Otherwise the keys are the index of the ``e_op``
            in the ``.expect`` list.

            The lists of expectation values returned are the *same* lists as
            those returned by ``.expect``.

        e_ops : dict
            A dictionary containing the supplied e_ops as ``ExpectOp`` instances.
            The keys of the dictionary are the same as for ``.e_data``.
            Each value is object where ``.e_ops[k](t, state)`` calculates the
            value of ``e_op`` ``k`` at time ``t`` and the given ``state``, and
            ``.e_ops[k].op`` is the original object supplied to create the
            ``e_op``.

        solver : str or None
            The name of the solver generating these results.

        stats : dict or None
            The stats generated by the solver while producing these results.

        options : dict
            The options for this result class.
        """

        times: list[float]
        states: list[Qobj]
        options: ResultOptions
        e_data: dict[Any, list[Any]]

        def __init__(
            self,
            e_ops: dict[Any, Qobj | QobjEvo | Callable[[float, Qobj], Any]],
            options: ResultOptions,
            *,
            solver: str = None,
            stats: dict[str, Any] = None,
            **kw,
        ):
            super().__init__(options, solver=solver, stats=stats)
            raw_ops = self._e_ops_to_dict(e_ops)
            self.e_data = {k: [] for k in raw_ops}
            self.e_ops = {}
            for k, op in raw_ops.items():
                f = self._e_op_func(op)
                self.e_ops[k] = ExpectOp(op, f, self.e_data[k].append)
                self.add_processor(self.e_ops[k]._store)

            self.times = []
            self.states = []
            self._final_state = None

            self._post_init(**kw)

        def _e_op_func(self, e_op):
            """
            Convert an e_op entry into a function, ``f(t, state)`` that returns
            the appropriate value (usually an expectation value).

            Sub-classes may override this function to calculate expectation values
            in different ways.
            """
            if isinstance(e_op, Qobj):
                return _QobjExpectEop(e_op)
            elif isinstance(e_op, QobjEvo):
                return e_op.expect
            elif callable(e_op):
                return e_op
            raise TypeError(f"{e_op!r} has unsupported type {type(e_op)!r}.")

        def _post_init(self):
            """
            Perform post __init__ initialisation. In particular, add state
            processors or pre-processors.

            Sub-class may override this. If the sub-class wishes to register the
            default processors for storing states, it should call this parent
            ``.post_init()`` method.

            Sub-class ``.post_init()`` implementation may take additional keyword
            arguments if required.
            """
            store_states = self.options["store_states"]
            store_states = store_states or (
                len(self.e_ops) == 0 and store_states is None
            )
            if store_states:
                self.add_processor(self._store_state, requires_copy=True)

            store_final_state = self.options["store_final_state"]
            if store_final_state and not store_states:
                self.add_processor(self._store_final_state, requires_copy=True)

        def _store_state(self, t, state):
            """Processor that stores a state in ``.states``."""
            self.states.append(state)

        def _store_final_state(self, t, state):
            """Processor that writes the state to ``._final_state``."""
            self._final_state = state

        def _pre_copy(self, state):
            """Return a copy of the state. Sub-classes may override this to
            copy a state in different manner or to skip making a copy
            altogether if a copy is not necessary.
            """
            return state.copy()

        def add(self, t, state):
            """
            Add a state to the results for the time ``t`` of the evolution.

            Adding a state calculates the expectation value of the state for
            each of the supplied ``e_ops`` and stores the result in ``.expect``.

            The state is recorded in ``.states`` and ``.final_state`` if specified
            by the supplied result options.

            Parameters
            ----------
            t : float
                The time of the added state.

            state : typically a :obj:`.Qobj`
                The state a time ``t``. Usually this is a :obj:`.Qobj` with
                suitable dimensions, but it sub-classes of result might support
                other forms of the state.

            Notes
            -----
            The expectation values, i.e. ``e_ops``, and states are recorded by
            the state processors (see ``.add_processor``).

            Additional processors may be added by sub-classes.
            """
            self.times.append(t)

            if (
                self._state_processors_require_copy
                and not self._integrator_return_copy
            ):
                state = self._pre_copy(state)

            for op in self._state_processors:
                op(t, state)

        def __repr__(self):
            lines = [
                f"<{self.__class__.__name__}",
                f"  Solver: {self.solver}",
            ]
            if self.stats:
                lines.append("  Solver stats:")
                lines.extend(f"    {k}: {v!r}" for k, v in self.stats.items())
            if self.times:
                lines.append(
                    f"  Time interval: [{self.times[0]}, {self.times[-1]}]"
                    f" ({len(self.times)} steps)"
                )
            lines.append(f"  Number of e_ops: {len(self.e_ops)}")
            if self.states:
                lines.append("  States saved.")
            elif self.final_state is not None:
                lines.append("  Final state saved.")
            else:
                lines.append("  State not saved.")
            lines.append(">")
            return "\n".join(lines)

        @property
        def expect(self) -> list[ArrayLike]:
            return [np.array(e_op) for e_op in self.e_data.values()]

        @property
        def final_state(self) -> Qobj:
            if self._final_state is not None:
                return self._final_state
            if self.states:
                return self.states[-1]
            return None