# Source code for qutip.solver.integrator.krylov

    from ..integrator import IntegratorException, Integrator
    import numpy as np
    from qutip.core import data as _data
    from scipy.optimize import root_scalar
    from ..sesolve import SESolver

    __all__ = ["IntegratorKrylov"]

    class IntegratorKrylov(Integrator):
        """
        Evolve the state vector ("psi0") finding an approximation for the time
        evolution operator of Hamiltonian ("H") by obtaining the projection of
        the time evolution operator on a set of small dimensional Krylov
        subspaces (m << dim(H)).
        """
        integrator_options = {
            'atol': 1e-7,
            'nsteps': 100,
            'min_step': 1e-5,
            'max_step': 1e5,
            'krylov_dim': 0,
            'sub_system_tol': 1e-7,
            'always_compute_step': False,
        }
        support_time_dependant = False
        supports_blackbox = False
        method = 'krylov'

        def _prepare(self):
            if not self.system.isconstant:
                raise ValueError("krylov method only support constant system.")
            self._max_step = -np.inf
            krylov_dim = self.options["krylov_dim"]
            if krylov_dim < 0 or krylov_dim > self.system.shape[0]:
                raise ValueError("The options 'krylov_dim', must be a positive "
                                 "integer smaller that the system size")

            if krylov_dim == 0:
                # TODO: krylov_dim, max_step and error (atol) are related by
                # err ~= exp(-krylov_dim / dt**(1~2))
                # We could ask for 2 and determine the third one.
                N = self.system.shape[0]
                krylov_dim = min(int((N + 100)**0.5), N-1)
                self.options["krylov_dim"]  = krylov_dim

            if not self.options["always_compute_step"]:
                from qutip import rand_ket
                N = self.system.shape[0]
                krylov_tridiag, krylov_basis = \
                    self._lanczos_algorithm(rand_ket(N).data)
                if (
                    krylov_tridiag.shape[0] < krylov_dim
                    or krylov_tridiag.shape[0] == N
                ):
                    self._max_step = np.inf
                else:
                    self._max_step = self._compute_max_step(krylov_tridiag,
                                                            krylov_basis)

        def _lanczos_algorithm(self, psi):
            """
            Computes a basis of the Krylov subspace for Hamiltonian 'H', a system
            state 'psi' and Krylov dimension 'krylov_dim'. The space is spanned
            by {psi, H psi, H^2 psi, ..., H^(krylov_dim) psi}.

            Parameters
            ------------
            psi: np.ndarray
                State used to calculate Krylov subspace.
            """
            krylov_dim = self.options['krylov_dim']
            H = (1j * self.system(0)).data

            v = []
            T_diag = np.zeros(krylov_dim + 1, dtype=complex)
            T_subdiag = np.zeros(krylov_dim + 1, dtype=complex)

            w_prime = _data.matmul(H, psi)
            T_diag[0] = _data.inner(w_prime, psi)
            v.append(psi)
            w = _data.add(w_prime, v[-1], -T_diag[0])
            T_subdiag[0] = _data.norm.l2(w)
            j = 0

            while j < krylov_dim and T_subdiag[j] > self.options['sub_system_tol']:
                j += 1
                v.append(_data.mul(w, 1 / T_subdiag[j-1]))
                w_prime = _data.matmul(H, v[-1])
                T_diag[j] = _data.inner(w_prime, v[-1])
                w = _data.add(w_prime, v[-1], -T_diag[j])
                w = _data.add(w, v[-2], -T_subdiag[j-1])
                T_subdiag[j] = _data.norm.l2(w)

            krylov_tridiag = _data.diag["dense"](
                [T_subdiag[:j], T_diag[:j+1], T_subdiag[:j]],
                [-1, 0, 1]
            )
            krylov_basis = _data.Dense(np.hstack([psi.to_array() for psi in v]))

            return krylov_tridiag, krylov_basis

        def _compute_krylov_set(self, krylov_tridiag, krylov_basis):
            """
            Compute the eigen energies, basis transformation operator (U) and e0.
            """
            eigenvalues, eigenvectors = _data.eigs(krylov_tridiag, True)
            N = eigenvalues.shape[0]
            U = _data.matmul(krylov_basis, eigenvectors)
            e0 = eigenvectors.adjoint() @ _data.one_element_dense((N, 1), (0, 0), 1.0)
            return eigenvalues, 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, krylov_basis, krylov_state=None):
            """
            Compute the maximum step length to stay under the desired tolerance.
            """
            if not krylov_state:
                krylov_state = self._compute_krylov_set(krylov_tridiag, krylov_basis)

            small_tridiag = _data.Dense(krylov_tridiag.as_ndarray()[:-1, :-1])
            small_basis = _data.Dense(krylov_basis.as_ndarray()[:, :-1])
            reduced_state = self._compute_krylov_set(small_tridiag, small_basis)

            def krylov_error(t):
                # we divide by atol and take the log so that the error returned is 0
                # at atol, which is convenient for calling root_scalar with.
                return np.log(_data.norm.l2(
                    self._compute_psi(t, *krylov_state) -
                    self._compute_psi(t, *reduced_state)
                ) / self.options["atol"])

            # Under 0 will cause an infinite loop in the while loop bellow.
            dt = max(self.options["min_step"], 1e-14)
            max_step = max(self.options["max_step"], dt)
            err = krylov_error(dt)
            if err > 0:
                raise ValueError(
                    f"With the krylov dim of {self.options['krylov_dim']}, the "
                    f"error with the minimum step {dt} is {err}, higher than the "
                    f"desired tolerance of {self.options['atol']}."
                )

            while krylov_error(dt * 10) < 0 and dt < max_step:
                dt *= 10

            if dt > max_step:
                return max_step

            sol = root_scalar(f=krylov_error, bracket=[dt, dt * 10],
                              method="brentq", xtol=self.options['atol'])
            if sol.converged:
                return sol.root
            else:
                return dt

        def set_state(self, t, state0):
            self._t_0 = t
            krylov_tridiag, krylov_basis = self._lanczos_algorithm(state0)
            self._krylov_state = self._compute_krylov_set(krylov_tridiag, krylov_basis)

            if (
                krylov_tridiag.shape[0] <= self.options['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, krylov_basis, self._krylov_state,
                )

        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(f"Maximum number of integration steps ({self.options['nsteps']}) exceeded")
                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
            out = self._compute_psi(delta_t, *self._krylov_state)
            return t, out

        @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 step size.

            krylov_dim: int, default: 0
                Dimension of Krylov approximation subspaces used for the time
                evolution approximation. If the defaut 0 is given, the dimension is calculated
                from the system size N, using `min(int((N + 100)**0.5), N-1)`.

            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')