"""Calculate the reservoir energy change from HOPS data."""

import qutip
import numpy as np
import scipy.misc

###############################################################################
#                              Single Trajectory                              #
###############################################################################


def flow_trajectory_coupling(ψ_0, ψ_1, L, G, W, bcf_scale=1, g=None, norm=None):
    """Calculates the $\langle L^\dagger \dot{B} + c.c.$ part of the
    energy flow for a single trajectory.

    :param ψ_0: the trajectory ``(time, dim-state)``
    :param ψ_1: the first HOPS hierarchy states ``(time, hierarchy-width * dim-state)``
    :param L: the coupling operator as system matrix
    :param G: the coupling factors in the exponential expansion of the BCF
    :param W: the exponents in the exponential expansion of the BCF
    :param bcf_scale: BCF scale factor
    :param g: individual scaling factors for the hierarchy states
    :param norm: the norm of the state, may be ``None`` for linear HOPS

    :returns: the value of the flow for each time step
    """

    t_steps, dim = ψ_0.shape
    hierarchy_width = ψ_1.shape[1] / dim

    ψ_coup = np.array((L @ ψ_0.T).T)

    # here we apply the prefactors to each hierarchy state
    ψ_hops = (1j * W * G * bcf_scale)[None, :, None] * ψ_1.reshape(
        t_steps, hierarchy_width, dim
    )

    # if there are per-term scaling factors
    if len(g) == hierarchy_width:
        ψ_hops /= g[None, :, None]

    flow = 2 * np.sum(ψ_coup.conj()[:, None, :] * ψ_hops, axis=(1, 2)).real

    # optionally normalize
    if norm:
        flow /= norm

    return flow


def flow_trajectory_therm(ψ_0, L, ξ, ξ_coeff, τ, norm=None, dx=1e-3, order=3):
    """Calculates the $\langle L^\dagger \dot{\xi} + c.c.$ part of the
    energy flow for a single trajectory.

    :param ψ_0: the trajectory ``(time, dim-state)``
    :param L: the coupling operator as system matrix
    :param ξ: the thermal stochastic process
    :param ξ_coeff: the coefficients of the realization of the
                    thermal stochastic process
    :param τ: the time points of the trajectory
    :param norm: the norm of the state, may be ``None`` for linear HOPS
    :param dx: step size for the calculation of the derivative of ξ
    :param order: order the calculation of the derivative of ξ, must be odd,
                  see :py:`scipy.misc.derivative`

    :returns: the value of the flow for each time step
    """

    ξ.new_process(ξ_coeff)
    ξ_dot = scipy.misc.derivative(ξ, τ, dx=dx, order=order)
    ψ_coup = np.array((L @ ψ_0.T).T) # TODO: prevent calculating this twice

    flow = 2 * (np.sum(ψ_coup.conj() * ψ_0)) * ξ_dot

    if norm:
        flow /= norm

    return flow