import numpy as np
from qutip.core import data as _data
from ..integrator import IntegratorException, Integrator
from ..sesolve import SESolver
from ..mesolve import MESolver


__all__ = ["IntegratorKrylov"]


class IntegratorKrylov(Integrator):
    """
    Evolve the state ("rho0") finding an approximation for the time evolution
    operator of a Hamiltonian ("H") by obtaining the projection on a set of
    small dimensional Krylov subspaces (m <= dim(H)). The construction of this
    subspace is performed by the Lanczos, fully-reorthogonalized Lanczos or
    Arnoldi algorithm.
    """
    integrator_options = {
        'atol': 1e-7,
        'nsteps': 100,
        'max_step': 1e5,
        'min_step': 1e-5,
        'always_compute_step': False,
        'krylov_dim': 0,
        'sub_system_tol': 1e-7,
        'algorithm': 'auto',
    }
    support_time_dependant = False
    supports_blackbox = False
    method = 'krylov'

    def _prepare(self):
        if not self.system.isconstant:
            raise ValueError("Krylov method only supports constant systems.")

        if self.options["krylov_dim"] <= 0:
            raise ValueError("The option 'krylov_dim', must be an integer "
                             "greater than zero.")

        self._max_step = -np.inf
        self._krylov_dim = self.options["krylov_dim"]
        self._hermitian = (1j*self.system(0)).isherm

        if self.options['algorithm'] == 'auto':
            if self._hermitian:
                self._algorithm = self._lanczos_full_reorth_algorithm
            else:
                self._algorithm = self._arnoldi_algorithm
        elif self.options['algorithm'] == 'arnoldi':
            self._algorithm = self._arnoldi_algorithm
        elif self.options['algorithm'] == 'lanczos_fro':
            self._algorithm = self._lanczos_full_reorth_algorithm
        elif self.options['algorithm'] == 'lanczos':
            self._algorithm = self._lanczos_algorithm
        else:
            raise ValueError("The requested algorithm "
                             f"{self.options['algorithm']} "
                             "for Krylov space construction is not available. "
                             "Possible options are: \'lanczos\', "
                             "\'lanczos_fro\', \'arnoldi\'.")

        if not self._hermitian and self._algorithm != self._arnoldi_algorithm:
            # Arnoldi is the only algorithm for open systems in QuTiP atm
            raise ValueError(f"The requested Krylov algorithm "
                             f"{self.options['algorithm']} "
                             "is not supported for non-Hermitian systems.")

    def _lanczos_algorithm(self, psi):
        return self._lanczos_core(psi, max_orthog_steps=2)

    def _lanczos_full_reorth_algorithm(self, psi):
        return self._lanczos_core(psi)

    def _lanczos_core(self, psi, max_orthog_steps=0):
        """
        Computes a basis of the Krylov subspace for the time independent
        Hamiltonian 'H', a system state 'psi' and Krylov dimension 'krylov_dim'
        using the Lanczos algorithm with a given orthogonalization function.
        The space is spanned by
        {psi, H psi, H^2 psi, ..., H^(krylov_dim - 1) psi}.

        Parameters
        ------------
        psi: np.ndarray
            State used to calculate Krylov subspace (= first basis state).
        max_orthog_steps: int
            Max. number of previous basis vectors to reorthogonalize against.

        Returns
        ------------
        krylov_trid: np.ndarray
            The tridiagonal matrix of the Krylov subspace.
        krylov_basis: np.ndarray
            The basis vectors of the Krylov subspace.
        """
        krylov_dim = self._krylov_dim
        H = (1j * self.system(0)).data
        p0 = _data.inner(psi, psi)  # purity
        sp0 = np.sqrt(p0)

        diag = np.zeros(krylov_dim, dtype=complex)
        subdiag = np.zeros(krylov_dim, dtype=complex)
        Q = [psi]

        k = 1
        v = _data.matmul(H, Q[-1])
        diag[0] = _data.inner(Q[-1], v) / p0
        v = _data.add(v, Q[-1], -diag[0])
        subdiag[0] = _data.norm.l2(v) / sp0
        while k < krylov_dim and subdiag[k-1] > self.options['sub_system_tol']:
            Q.append(_data.mul(v, 1 / subdiag[k-1]))
            v = _data.matmul(H, Q[-1])
            k += 1
            v, ol = self._orthogonalize(Q, v, p0, steps=max_orthog_steps)
            diag[k-1] = ol
            subdiag[k-1] = _data.norm.l2(v) / sp0

        krylov_trid = _data.diag["dense"](
            [subdiag[:k-1], diag[:k], subdiag[:k-1]],
            [-1, 0, 1]
        )
        krylov_basis = _data.Dense(np.hstack([p.to_array() for p in Q]))

        return krylov_trid, krylov_basis

    def _orthogonalize(self, Q, v, p0, steps=0):
        """
        Orthogonalizes a new vector `v` against the previous `max_orthog_steps`
        number of Krylov basis vectors in the list `Q`.

        Parameters
        ------------
        Q: list of np.ndarray
            The list of previous Krylov basis vectors.
        v: np.ndarray
            The new vector to orthogonalize.
        p0: float
            The purity of the initial state.
        steps: int, default: 0
            The number of previous basis vectors to reorthogonalize against.
            The default `0` will reorthogonalize w.r.t all.

        Returns
        ------------
        v: np.ndarray
            The orthogonalized vector = new Krylov basis vector.
        ol: float
            The overlap of the orthogonalized vector with the last basis vector
            in `Q`. This is the new off diagonal element of the tridiagonal
            matrix.
        """
        ol = 0
        for q in Q[-steps:]:
            ol = _data.inner(q, v) / p0
            v = _data.add(v, q, -ol)
        return v, ol

    def _arnoldi_algorithm(self, psi):
        """
        Computes the Krylov subspace basis for a Hamiltonian 'H', a system
        state 'psi' and Krylov dimension 'krylov_dim' using the Arnoldi
        interation. This results in an upper Hessenberg matrix that in general
        is non-Hermitian. The space is spanned by
        {psi, H psi, H^2 psi, ..., H^(krylov_dim - 1) psi}.

        Parameters
        ------------
        psi: np.ndarray
            State used to calculation Krylov subspace (= first basis state).

        Returns
        ------------
        krylov_hesse: np.ndarray
            The upper triangular matrix of the Krylov subspace.
        krylov_basis: np.ndarray
            The basis vectors of the Krylov subspace.
        """
        krylov_dim = self._krylov_dim
        H = (1j * self.system(0)).data
        p0 = _data.inner(psi, psi)  # purity
        sp0 = np.sqrt(p0)

        h = np.zeros((krylov_dim + 1, krylov_dim), dtype=complex)
        Q = [psi]

        k = 1
        v = _data.matmul(H, Q[-1])
        h[0, 0] = _data.inner(Q[-1], v) / p0
        v = _data.add(v, Q[-1], -h[0, 0])
        h[1, 0] = _data.norm.l2(v) / sp0
        while k < krylov_dim and h[k, k-1] > self.options['sub_system_tol']:
            Q.append(_data.mul(v, 1 / h[k, k-1]))
            v = _data.matmul(H, Q[-1])
            k += 1
            for j in range(k):  # removes projections, create upper Hessenberg
                h[j, k-1] = _data.inner(Q[j], v) / p0
                v = _data.add(v, Q[j], -h[j, k-1])
            h[k, k-1] = _data.norm.l2(v) / sp0

        krylov_hesse = _data.Dense(h[:k, :k])
        krylov_basis = _data.Dense(
            np.hstack([p.to_array() for p in Q])
        )
        return krylov_hesse, krylov_basis

    def _compute_krylov_set(self, krylov_tridiag, krylov_basis):
        """
        Compute the eigen energies, basis transformation operator (U) and e0.
        """
        evals, evecs = _data.eigs(krylov_tridiag, self._hermitian)
        N = evals.shape[0]
        U = _data.matmul(krylov_basis, evecs)

        e0 = _data.one_element_dense((N, 1), (0, 0), 1.0)
        if self._hermitian:
            e0 = evecs.adjoint() @ e0
        else:
            e0 = _data.inv(evecs) @ e0

        return evals, U, e0

    def _compute_psi(self, dt, eigenvalues, U, e0):
        """
        Compute the state at time ``t``.
        """
        phases = _data.Dense(np.exp(-1j * dt * eigenvalues))
        aux = _data.multiply(phases, e0)
        return _data.matmul(U, aux)

    def _compute_max_step(self, krylov_tridiag):
        """
        Compute the maximum step length to stay under the desired tolerance.
        """
        n = krylov_tridiag.shape[0]
        facto = np.arange(2, n+1)
        bsprod = np.diag(krylov_tridiag.as_ndarray(), k=-1)

        lnum = np.log(self.options["atol"]) + np.sum(np.log(facto))
        lden = np.sum(np.log(bsprod))
        dt = np.exp(np.real((lnum - lden) / n))

        if dt < self.options["min_step"]:
            raise ValueError(
                f"With the krylov dimension of {self.options['krylov_dim']} "
                f"and desired tolerance of {self.options['atol']}, the "
                f"maximum possible time step size is {dt}. But is smaller "
                f"than the minimum desired time step size of "
                f"{self.options['min_step']}."
            )
        return min(dt, self.options["max_step"])

    def set_state(self, t, state0):
        self._t_0 = t

        krylov_tridiag, krylov_basis = self._algorithm(state0)
        self._krylov_state = \
            self._compute_krylov_set(krylov_tridiag, krylov_basis)

        if (
            krylov_tridiag.shape[0] < self._krylov_dim
            or krylov_tridiag.shape == self.system.shape
        ):
            # happy_breakdown
            self._max_step = np.inf
            return

        if (
            not np.isfinite(self._max_step)
            or self.options["always_compute_step"]
        ):
            self._max_step = self._compute_max_step(krylov_tridiag)

    def get_state(self, copy=True):
        return self._t_0, self._compute_psi(0, *self._krylov_state)

    def integrate(self, t, copy=True):
        step = 0
        while t > self._t_0 + self._max_step:
            # The approximation in only valid in the range t_0, t_0 + max step
            # If outside, advance the range
            step += 1
            if step >= self.options["nsteps"]:
                raise IntegratorException(
                    "Maximum number of integration steps "
                    f"({self.options['nsteps']}) exceeded. "
                    "Increase the number of steps or krylov dimension or "
                    "reduce tolerance."
                )
            new_psi = self._compute_psi(self._max_step, *self._krylov_state)
            self.set_state(self._t_0 + self._max_step, new_psi)

        delta_t = t - self._t_0
        return t, self._compute_psi(delta_t, *self._krylov_state)

    @property
    def options(self):
        """
        Supported options by krylov method:

        atol : float, default: 1e-7
            Absolute tolerance.

        nsteps : int, default: 100
            Max. number of internal steps/call.

        min_step, max_step : float, default: (1e-5, 1e5)
            Minimum and maximum time step size before the Krylov basis is
            recalculated.

        krylov_dim: int, default: 0
            Dimension of Krylov approximation subspaces used for the time
            evolution approximation.

        algorithm: str, default: "auto"
            Algorithm for Krylov space constructions. The default ``auto`` will
            choose ``lanczos_fro`` for Hermitian and ``arnoldi`` for
            non-Hermitian systems. Alternatively the standard ``lanczos`` can
            be set.

        sub_system_tol: float, default: 1e-7
            Tolerance to detect a happy breakdown. A happy breakdown occurs
            when the initial ket is in a subspace of the Hamiltonian smaller
            than ``krylov_dim``.

        always_compute_step: bool, default: False
            If True, the step length is computed each time a new Krylov
            subspace is computed. Otherwise it is computed only once when
            creating the integrator.
        """
        return self._options

    @options.setter
    def options(self, new_options):
        Integrator.options.fset(self, new_options)


SESolver.add_integrator(IntegratorKrylov, 'krylov')
MESolver.add_integrator(IntegratorKrylov, 'krylov')
