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
