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


def make_hops_step(
    η, H_s: np.ndarray, L: np.ndarray, W: complex, k_max: int
) -> t.Callable[float, np.ndarray]:
    """Creates the step function for 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
    """
    dim = H_s.shape[0]

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

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

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

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

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

        res = np.array((ψ_1 + ψ_2 + ψ_3)).reshape(((k_max + 1) * dim,), order="F")
        # print(res)
        # __import__("pdb").set_trace()
        return res

    return step


def integrate_hops_trajectory(
    η,
    H_s: np.ndarray,
    L: np.ndarray,
    W: complex,
    k_max: int,
    ψ_0: np.ndarray,
    τ_max: float,
    seed: t.Optional[int] = None,
    **kwargs
):
    dim = H_s.shape[0]
    ψ_0_ext = np.concatenate((ψ_0, np.zeros(k_max * dim))) + 0j
    η.new_process(seed=seed)
    step = make_hops_step(η, H_s, L, W, k_max)
    return scipy.integrate.solve_ivp(
        step, (0, τ_max), ψ_0_ext, vectorized=False, dense_output=True, **kwargs
    )


def integrate_hops_ensemble(
    η,
    H_s: np.ndarray,
    L: np.ndarray,
    W: complex,
    k_max: int,
    ψ_0: np.ndarray,
    τ_max: float,
    N: int,
    **kwargs
):
    dim = H_s.shape[0]
    τ = np.linspace(0, τ_max, 100)
    ρ = np.zeros((100, dim, dim), dtype="complex128")

    for _ in range(0, N):
        ψ = (
            integrate_hops_trajectory(
                η, H_s, L, W, k_max, ψ_0, τ_max, seed=None, **kwargs
            )
            .sol(τ)[0:2, :]
            .T
        )

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

    return τ, ρ / N