"""Calculate the energy flow into the bath for a simple gaussian
  Quantum Brownian Motion model.

  This is done analytically for a BCF that is a sum of exponentials.
"""

import numpy as np
import itertools
from dataclasses import dataclass, field, InitVar
from typing import Callable, Union
import numpy.typing as npt
import lmfit
from . import util
import functools
from numpy.polynomial import Polynomial


@dataclass
class BCF:
    r"""A parameter object to hold information about a BCF.

    The BCFs will be expanded into a sum of exponentials like
    :math:`\alpha(\tau) = \sum_k G_k \cdot \exp(-W_k\cdot\tau)`.
    You can either give the BCFs as parameter or the coefficients.
    If you give the BCFs, the fit will be performed automatically.

    :param t_max: the maximum simulation time
    :param function: the BCF as python function, will be set to the
        exponential expansion if the BCF coefficients are given.

    :param num_terms: the number of terms of the expansion of the
        BCF
    :param precision: the precision in the sampling for the fit,
        ``t_max/precision`` points will be used
    :param exponents: the exponential factors in the BCF expansion
    :param factors: the pre-factors in the BCF expansion

    :attribute approx:
    """

    #: the maximum simulation time
    t_max: float
    function: Union[Callable[[npt.ArrayLike], np.ndarray], None] = None
    num_terms: Union[int, None] = None
    resolution: Union[float, None] = None
    factors: Union[np.ndarray, None] = None
    exponents: Union[np.ndarray, None] = None

    #: the BCF as exponential expansion
    approx: Callable[[npt.ArrayLike], np.ndarray] = field(init=False)

    def __post_init__(self):
        if self.function is not None:
            if self.num_terms is None or self.resolution is None:
                raise ValueError(
                    "Either give the function, the number of terms and the resolution or the coefficients."
                )

            self.exponents, self.factors = util.fit_α(
                self.function,
                self.num_terms,
                self.t_max,
                int(self.t_max / self.resolution),
            )

        else:
            if self.factors is None or self.exponents is None:
                raise ValueError(
                    "Either give the function and number of terms or the coefficients."
                )

            if self.factors.size != self.exponents.size:
                raise ValueError(
                    "Factors and exponents have to have the same dimension."
                )

            self.num_terms = self.factors.size

        self.approx = functools.partial(util.α_apprx, G=self.factors, W=self.exponents)

        if self.function is None:
            self.function = self.approx


@dataclass
class SystemParams:
    r"""A parameter object to hold information about the physical
    system and the global parameters for the algorithm.

    :param Ω: the system hamiltonian energy scale, :math:`H_s =
              \Omega (p^2 + q^2)`
    :param η: the coupling strength (in essence a prefactor to the BCF)
    :param α: the BCF


    :attr t_max: the maximum simulation time, will be copied from :attr:`α`
    :attr W: the exponential factors in the BCF expansion
    :attr G: the pre-factors in the BCF expansion
    """

    Ω: float
    η: float
    α: BCF

    t_max: float = field(init=False)

    # BCF coefficients
    W: np.ndarray = field(init=False)
    G: np.ndarray = field(init=False)

    def __post_init__(self):
        self.t_max = self.α.t_max

        self.W = self.α.exponents
        self.G = self.α.factors * (self.η ** 2)


def construct_polynomials(sys: SystemParams) -> tuple[Polynomial, Polynomial]:
    r"""
    Construct the polynomials required to find the coefficients
    and exponents of the solution.

    :param sys: a parameter object with the system parameters
    :returns: a list of polynomials

        - :math:`f_0(z) = \prod_k (z-z_k) (z-z_k^{\ast})`
        - :math:`p(z) = p_1(z) + \sum_n q_n(z)` where
            - :math:`p_1(z) = (z^2 + \Omega^2)\prod_k (z-z_k)(z-z_k^{\ast})`
            - :math:`q_n(z)=\Omega f_n(z) \prod_{k\neq n}(z-z_k) (z-z_k^{\ast})`

    """
    # we begin by calculating all the constants we'll need
    φ = -np.angle(sys.G)
    G_abs = np.abs(sys.G)
    γ, δ = sys.W.real, sys.W.imag
    s, c = np.sin(φ), np.cos(φ)

    # the roots of -G ((z + γ)s + δc)
    roots_f = -(γ + δ * c / s)

    # the roots of the denominator of the laplace transform of α
    roots_z = -sys.W

    # now the polynomials
    f_0 = util.poly_real(
        Polynomial.fromroots(np.concatenate((roots_z, roots_z.conj())))
    )
    p_1 = Polynomial([sys.Ω ** 2, 0, 1]) * f_0

    q = [
        -G_c
        * sys.Ω
        * s_c
        * util.poly_real(
            Polynomial.fromroots(
                np.concatenate(
                    (
                        [root_f],
                        util.except_element(roots_z, i),
                        util.except_element(roots_z, i).conj(),
                    )
                )
            )
        )
        for root_f, G_c, γ_c, δ_c, s_c, c_c, i in zip(
            roots_f, G_abs, γ, δ, s, c, range(len(c))
        )
    ]

    p = p_1 + sum(q)

    return f_0, p


def calculate_coefficients(sys: SystemParams) -> tuple[np.ndarray, np.ndarray]:
    r"""
    Calculates the coefficients required to construct the
    propagator matrix :math:`G`.

    :param sys: a parameter object with the system parameters
    :returns: the exponents and the residuals which play a role as pre-factors
    """

    f_0, p = construct_polynomials(sys)

    master_roots = p.roots()

    if np.unique(master_roots).size != master_roots.size:
        raise RuntimeError(
            """The roots of the polynomial are not unique.
You can try to alter the number of terms in the expansion of the BCF."""
        )

    resiquals = f_0(master_roots) / p.deriv()(master_roots)
    return master_roots, resiquals


def _expand_t(f):
    def wrapped(self, t):
        t = np.expand_dims(np.asarray(t), axis=0)
        return f(self, t)

    return wrapped


class Propagator:
    """The propagator matrix :math:`G` for the system.

    You can get inidividual matrix elements as functions of time
    through indexing this object.

    Calling it with a time argument returns the whole matrix.

    :param sys: a parameter object with the system parameters
    """

    def __init__(self, sys: SystemParams):
        self.params = sys
        self._roots, self._res = calculate_coefficients(sys)
        self._roots_exp = np.expand_dims(self._roots, axis=1)
        self._res_exp = np.expand_dims(self._res, axis=1)
        self._Ω = sys.Ω

        self._res_times_roots = self._res_exp * self._roots_exp
        self._res_times_roots_squared = self._res_times_roots * self._roots_exp
        self._elements = np.array([[self.el_11, self.el_12], [self.el_21, self.el_22]])
        pass

    @_expand_t
    def el_11(self, t):
        return (self._res_times_roots * np.exp(t * self._roots_exp)).real.sum(axis=0)

    @_expand_t
    def el_12(self, t):
        return self._Ω * (self._res_exp * np.exp(t * self._roots_exp)).real.sum(axis=0)

    @_expand_t
    def el_21(self, t):
        return (self._res_times_roots_squared * np.exp(t * self._roots_exp)).real.sum(
            axis=0
        ) / self._Ω

    def el_22(self, t):
        return self.el_11(t)

    def __getitem__(self, indices) -> Callable[[npt.ArrayLike], np.ndarray]:
        """Get an individual matrix elment as function."""
        return self._elements[indices]

    @_expand_t
    def __call__(self, t) -> np.ndarray:
        """Get the propagator as array of shape ``(time, 2, 2)``."""

        g_12 = self._res_exp * np.exp(t * self._roots_exp)
        diag = self._roots_exp * g_12
        g_21 = self._roots_exp * diag / self._Ω

        return (
            np.array([[diag, g_12 * self._Ω], [g_21, diag]])
            .real.sum(axis=2)
            .swapaxes(0, 2)
            .swapaxes(1, 2)
        )


@dataclass
class FlowParams:
    r"""
    A parameter object to hold the global parameters for the flow algorithm.

    :param system: the system parameters, see :any:`SystemParams`
    :param α_0_dot: the zero temperature BCF time derivative
    :param n: the excitation number of the initial state of the system :math:`|n\rangle`

    :attr L: the exponential factors in the BCF derivative
             expansion :math:`\dot{\alpha}_0=\sum_k P_k e^{-L_k \cdot t}`

    :attr P: the pre-factors in the BCF derivative expansion
    :attr q_s_0: the expectation value :math:`\langle q(0)^2\rangle`
    :attr p_s_0: the expectation value :math:`\langle q(0)^2\rangle`
    :attr qp_0: the expectation value :math:`\langle q(0)p(0)\rangle`
    :attr prop: the propagator matrix :math:`G(t)`, see :any:`Propagator`
    :attr B_coef: the pre-factors of the exponential sum of :math:`B=G_{12}=\sum_k B_k e^{-C_k \cdot t}`

                    - note the minus sign
    :attr C: the exponents of the exponential sum of ``B``
    :attr A_coef: the pre-factors of the exponential sum of :math:`A=G_{11}=\sum_k A_k e^{-C_k \cdot t}`
    """

    system: InitVar[SystemParams]
    α_0_dot: BCF
    n: int

    # BCF derivative coefficients
    L: np.ndarray = field(init=False)
    P: np.ndarray = field(init=False)

    # initial state
    q_s_0: float = field(init=False)
    p_s_0: float = field(init=False)
    qp_0: complex = field(init=False)

    # propagator
    prop: Propagator = field(init=False)
    B_coef: np.ndarray = field(init=False)
    C: np.ndarray = field(init=False)
    A_coef: np.ndarray = field(init=False)

    def __post_init__(self, system):
        self.L = self.α_0_dot.exponents
        self.P = self.α_0_dot.factors * (system.η ** 2)

        self.q_s_0 = 1 + 2 * self.n
        self.p_s_0 = 1 + 2 * self.n
        self.qp_0 = 1j

        self.C, self.B_coef = calculate_coefficients(system)
        self.A_coef = self.B_coef * self.C

        self.B_coef = system.Ω * self.B_coef
        self.C = -self.C  # mind the extra -

        self.prop = Propagator(system)


def get_flow(
    sys: SystemParams, flow: FlowParams
) -> Callable[[npt.ArrayLike], np.ndarray]:
    r"""Get the flow :math:`\frac{d}{dt}\langle H_B\rangle` as a function of time."""

    # TODO: Think about the too bloated ``FlowParams`` class and this ugly function
    A = flow.prop[0, 0]
    B = flow.prop[0, 1]
    Bk = flow.B_coef  #
    Ak = flow.A_coef  #
    Pk = flow.P  #
    Lk = flow.L  #
    Ck = flow.C  #
    Wk = sys.W  #
    Gk = sys.G  #
    Γ1 = (Bk[:, None] * Pk[None, :]) / (Lk[None, :] - Ck[:, None])
    Γ2 = (Bk[:, None] * Gk[None, :]) / (Ck[:, None] - Wk[None, :])
    Γ3 = (Bk[:, None] * Gk.conj()[None, :]) / (Ck[:, None] + Wk.conj()[None, :])
    ΓA = (Ak[:, None] * Pk[None, :]) / (Lk[None, :] - Ck[:, None])

    def A_conv(t):
        t = np.asarray(t)
        result = np.zeros_like(t, dtype="complex128")

        # print(np.sort(Ck.imag.flatten()))

        for (n, m) in itertools.product(range(len(Ak)), range(len(Pk))):
            result += ΓA[n, m] * (np.exp(-Ck[n] * t) - np.exp(-Lk[m] * t))

        return result

    def B_conv(t):
        t = np.asarray(t)
        result = np.zeros_like(t, dtype="complex128")
        for (n, m) in itertools.product(range(len(Bk)), range(len(Pk))):
            result += Γ1[n, m] * (np.exp(-Ck[n] * t) - np.exp(-Lk[m] * t))

        return result

    def flow_conv_free(t):
        a, b = A(t), B(t)
        ac, bc = A_conv(t), B_conv(t)
        # print(ac, bc)
        return (
            -1
            / 2
            * (
                flow.q_s_0 * a * ac
                + flow.p_s_0 * b * bc
                + flow.qp_0 * a * bc
                + flow.qp_0.conjugate() * b * ac
            ).imag
        )

    def conv_part(t):
        t = np.asarray(t)
        result = np.zeros_like(t, dtype="float")
        for (m, k, n, l) in itertools.product(
            range(len(Bk)), range(len(Pk)), range(len(Bk)), range(len(Gk))
        ):
            g_1_2 = (
                Γ1[m, k]
                * Γ2[n, l]
                * (
                    (1 - np.exp(-(Ck[m] + Wk[l]) * t)) / (Ck[m] + Wk[l])
                    + (np.exp(-(Ck[m] + Ck[n]) * t) - 1) / (Ck[m] + Ck[n])
                    + (np.exp(-(Lk[k] + Wk[l]) * t) - 1) / (Lk[k] + Wk[l])
                    + (1 - np.exp(-(Lk[k] + Ck[n]) * t)) / (Lk[k] + Ck[n])
                )
            )

            g_1_3 = (
                Γ1[m, k]
                * Γ3[n, l]
                * (
                    (1 - np.exp(-(Ck[m] + Ck[n]) * t)) / (Ck[m] + Ck[n])
                    - (1 - np.exp(-(Lk[k] + Ck[n]) * t)) / (Lk[k] + Ck[n])
                    - (
                        np.exp(-(Ck[n] + Wk[l].conj() * t))
                        - np.exp(-(Ck[m] + Ck[n]) * t)
                    )
                    / (Ck[m] - Wk[l].conj())
                    + (
                        np.exp(-(Ck[n] + Wk[l].conj() * t))
                        - np.exp(-(Lk[k] + Ck[n]) * t)
                    )
                    / (Lk[k] - Wk[l].conj())
                )
            )

            result += -1 / 2 * (g_1_2.imag + g_1_3.imag)

        return result

    def final_flow(t: npt.ArrayLike) -> np.ndarray:
        return flow_conv_free(t) + conv_part(t)

    return final_flow