from qibo import K
from qibo.abstractions import hamiltonians
from qibo.config import raise_error


class BaseSolver:
    """Basic solver that should be inherited by all solvers.

    Args:
        dt (float): Time step size.
        hamiltonian (:class:`qibo.abstractions.hamiltonians.Hamiltonian`): Hamiltonian object
            that the state evolves under.
    """

    def __init__(self, dt, hamiltonian):
        self.dt = dt
        if issubclass(type(hamiltonian), hamiltonians.HAMILTONIAN_TYPES):
            self.hamiltonian = lambda t: hamiltonian
        else:
            self.hamiltonian = hamiltonian
        self.t = 0

    @property
    def t(self):
        """Solver's current time."""
        return self._t

    @t.setter
    def t(self, new_t):
        """Updates solver's current time."""
        self._t = new_t
        self.current_hamiltonian = self.hamiltonian(self.t)

    def __call__(self, state): # pragma: no cover
        # abstract method
        raise_error(NotImplementedError)


class TrotterizedExponential(BaseSolver):
    """Solver that uses Trotterized exponentials.

    Created automatically from the :class:`qibo.solvers.Exponential` if the
    given Hamiltonian object is a
    :class:`qibo.abstractions.hamiltonians.TrotterHamiltonian`.
    """

    def __call__(self, state):
        circuit = self.hamiltonian(self.t).circuit(self.dt)
        self.t += self.dt
        return circuit(state)


class Exponential(BaseSolver):
    """Solver that uses the matrix exponential of the Hamiltonian:

    .. math::
        U(t) = e^{-i H(t) \\delta t}

    Calculates the evolution operator in every step and thus is compatible with
    time-dependent Hamiltonians.
    """

    def __new__(cls, dt, hamiltonian):
        if issubclass(type(hamiltonian), hamiltonians.HAMILTONIAN_TYPES):
            h0 = hamiltonian
        else:
            h0 = hamiltonian(0)
        if isinstance(h0, hamiltonians.TrotterHamiltonian):
            return TrotterizedExponential(dt, hamiltonian)
        else:
            return super(Exponential, cls).__new__(cls)

    def __call__(self, state):
        propagator = self.current_hamiltonian.exp(self.dt)
        self.t += self.dt
        return K.matmul(propagator, state[:, K.newaxis])[:, 0]


class RungeKutta4(BaseSolver):
    """Solver based on the 4th order Runge-Kutta method."""

    def __call__(self, state):
        ham1 = self.current_hamiltonian
        ham2 = self.hamiltonian(self.t + self.dt / 2.0)
        ham3 = self.hamiltonian(self.t + self.dt)
        k1 = ham1 @ state
        k2 = ham2 @ (state + self.dt * k1 / 2.0)
        k3 = ham2 @ (state + self.dt * k2 / 2.0)
        k4 = ham3 @ (state + self.dt * k3)
        self.t += self.dt
        return (state - 1j * self.dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6.0)


class RungeKutta45(BaseSolver):
    """Solver based on the 5th order Runge-Kutta method."""

    def __call__(self, state):
        ham1 = self.current_hamiltonian
        ham2 = self.hamiltonian(self.t + self.dt / 4.0)
        ham3 = self.hamiltonian(self.t + 3 * self.dt / 8.0)
        ham4 = self.hamiltonian(self.t + 12 * self.dt / 13.0)
        ham5 = self.hamiltonian(self.t + self.dt)
        ham6 = self.hamiltonian(self.t + self.dt / 2.0)
        k1 = ham1 @ state
        k2 = ham2 @ (state + self.dt * k1 / 4.0)
        k3 = ham3 @ (state + self.dt * (3 * k1 + 9 * k2) / 32.0)
        k4 = ham4 @ (state + self.dt * (1932 * k1 -
                                        7200 * k2 + 7296 * k3) / 2197.0)
        k5 = ham5 @ (state + self.dt * (439 * k1 / 216.0 - 8 *
                                        k2 + 3680 * k3 / 513.0 - 845 * k4 / 4104.0))
        k6 = ham6 @ (state + self.dt * (-8 * k1 / 27.0 + 2 * k2 -
                                        3544 * k3 / 2565 + 1859 * k4 / 4104 - 11 * k5 / 40.0))
        self.t += self.dt
        return (state - 1j * self.dt * (16 * k1 / 135.0 + 6656 * k3 / 12825.0 + 28561 * k4 / 56430.0 -
                                        9 * k5 / 50.0 + 2 * k6 / 55.0))


factory = {
    "exp": Exponential,
    "rk4": RungeKutta4,
    "rk45": RungeKutta45
}
