"""A pedadogical exercise implementation of HOPS."""
import numpy as np
import numpy.typing as npt
import typing as t
import pdb
import scipy


class Hops:
    def __init__(self, η, H_s, L, W, k_max, seed=None, solver_args=dict()):
        """Implements the integration of hops with
        exactly one exponential term.

        :param η: random process from `stocproc`
        :param H_s: the system hamiltonian
        :param L: the interaction operator
        :param W: the exponent of the BCF
        :param k_max: the depth of the hirarchy
        :returns: the step function for the integration of hops
        """
        self.η = η
        self.H_s = H_s
        self.L = L
        self.W = W
        self.k_max = k_max
        self.dim = self.H_s.shape[0]
        self.step = self._make_hops_step()
        self.seed = seed
        self._solver_args = solver_args


    def _make_hops_step(self):
        H_si = -1j * self.H_s
        B = -np.conj(self.L).T
        K = np.diag(np.arange(0, self.k_max + 1))

        def step(t, ψ):
            ψ = ψ.reshape((self.dim, self.k_max + 1), order="F")

            # 1. Apply system H and tho constant contributions
            ψ_1 = H_si @ ψ + self.W * (ψ @ K) + self.L @ (ψ * self.η(t))

            # 2. Now the shifted orders, we set the truncator to zero
            zeros = np.zeros((1, self.dim)).T
            ψ_ext = np.hstack((zeros, ψ, zeros))

            ψ_2 = B @ ψ_ext[:, 0:-2]
            ψ_3 = self.L @ (ψ_ext[:, 2:] @ K)

            res = np.array((ψ_1 + ψ_2 + ψ_3)).reshape(
                ((self.k_max + 1) * self.dim,), order="F"
            )

            return res

        return step

    def integrate_hops_trajectory(self, ψ_0, τ_max, seed=None):
        seed = self.seed if seed is None else seed
        ψ_0_ext = np.concatenate((ψ_0, np.zeros(self.k_max * self.dim))) + 0j
        self.η.new_process(seed=seed)
        return scipy.integrate.solve_ivp(
            self.step,
            (0, τ_max),
            ψ_0_ext,
            vectorized=False,
            dense_output=True,
            **self._solver_args
        )

    def integrate_hops_ensemble(self, ψ_0, τ, N):
        τ_max = np.max(τ)
        ρ = np.zeros((len(τ), self.dim, self.dim), dtype="complex128")

        for i in range(0, N):
            ψ = (
                self.integrate_hops_trajectory(
                    ψ_0,
                    τ_max,
                    seed=(None if self.seed is None else self.seed + i)
                )
                .sol(τ)[0:2, :]
                .T
            )

            ρ += ψ[:, :, np.newaxis] * ψ.conj()[:, np.newaxis, :]

        return τ, ρ / N