"""
Implementation of the HFSModel class, currently only supplied with a Voigt profile.

.. moduleauthor:: Wouter Gins <wouter.gins@kuleuven.be>
"""
from __future__ import annotations

from typing import Tuple

import numpy as np
import uncertainties as unc
from numpy.typing import ArrayLike
from scipy.special import voigt_profile, wofz
from sympy.physics.wigner import wigner_3j, wigner_6j

from satlas2.core import Model, Parameter

__all__ = ['HFS']

sqrt2 = 2**0.5
sqrt2log2t2 = 2 * np.sqrt(2 * np.log(2))
log2 = np.log(2)


class HFS(Model):
    """Initializes a hyperfine spectrum Model with the given hyperfine parameters.

    Parameters
    ----------
    I : float
        Integer or half-integer value of the nuclear spin
    J : ArrayLike
        A sequence of 2 spins, respectively the J value of the lower state
        and the J value of the higher state
    A : ArrayLike, optional
        A sequence of 2 A values, respectively for the lower and the higher state, by default [0, 0]
    B : ArrayLike, optional
        A sequence of 2 B values, respectively for the lower and the higher state, by default [0, 0]
    C : ArrayLike, optional
        A sequence of 2 C values, respectively for the lower and the higher state, by default [0, 0]
    df : float, optional
        The centroid of the spectrum, by default 0
    fwhmg : float, optional
        The Gaussian FWHM of the Voigt profile, by default 50
    fwhml : float, optional
        The Lorentzian FWHM of the Voigt profile, by default 50
    name : str, optional
        Name of the model, by default 'HFS'
    N : int, optional
        Number of sidepeaks to be generated, by default None
    offset : float, optional
        Offset in units of x for the sidepeak, by default 0
    poisson : float, optional
        The poisson factor for the sidepeaks, by default 0
    scale : float, optional
        The amplitude of the entire spectrum, by default 1.0
    racah : bool, optional
        Use individual amplitudes are setting the Racah intensities, by default True
    prefunc : callable, optional
        Transformation to be applied on the input before evaluation, by default None
        """
    def __init__(self,
                 I: float,
                 J: ArrayLike,
                 A: ArrayLike = [0, 0],
                 B: ArrayLike = [0, 0],
                 C: ArrayLike = [0, 0],
                 df: float = 0,
                 fwhmg: float = 50,
                 fwhml: float = 50,
                 name: str = 'HFS',
                 N: int = None,
                 offset: float = 0,
                 poisson: float = 0,
                 scale: float = 1.0,
                 racah: bool = True,
                 prefunc: callable = None):
        super().__init__(name, prefunc=prefunc)
        J1, J2 = J
        lower_F = np.arange(abs(I - J1), I + J1 + 1, 1)
        upper_F = np.arange(abs(I - J2), I + J2 + 1, 1)

        self.lines = []
        self.intensities = {}
        self.scaling_Al = {}
        self.scaling_Bl = {}
        self.scaling_Cl = {}
        self.scaling_Au = {}
        self.scaling_Bu = {}
        self.scaling_Cu = {}

        for i, F1 in enumerate(lower_F):
            for j, F2 in enumerate(upper_F):
                if abs(F2 - F1) <= 1 and not F2 == F1 == 0.0:
                    if F1 % 1 == 0:
                        F1_str = '{:.0f}'.format(F1)
                    else:
                        F1_str = '{:.0f}_2'.format(2 * F1)

                    if F2 % 1 == 0:
                        F2_str = '{:.0f}'.format(F2)
                    else:
                        F2_str = '{:.0f}_2'.format(2 * F2)

                    line = '{}to{}'.format(F1_str, F2_str)
                    self.lines.append(line)

                    C1, D1, E1 = self.calcShift(I, J1, F1)
                    C2, D2, E2 = self.calcShift(I, J2, F2)

                    self.scaling_Al[line] = C1
                    self.scaling_Bl[line] = D1
                    self.scaling_Cl[line] = E1
                    self.scaling_Au[line] = C2
                    self.scaling_Bu[line] = D2
                    self.scaling_Cu[line] = E2

                    intens = float((2 * F1 + 1) * (2 * F2 + 1) * \
                         wigner_6j(J2, F2, I, F1, J1, 1.0) ** 2) # DO NOT REMOVE CAST TO FLOAT!!!
                    self.intensities['Amp' + line] = Parameter(value=intens,
                                                               min=0,
                                                               vary=not racah)

        norm = max([p.value for p in self.intensities.values()])
        for n, v in self.intensities.items():
            v.value /= norm

        pars = {
            'centroid': Parameter(value=df),
            'Al': Parameter(value=A[0]),
            'Au': Parameter(value=A[1]),
            'Bl': Parameter(value=B[0]),
            'Bu': Parameter(value=B[1]),
            'Cl': Parameter(value=C[0]),
            'Cu': Parameter(value=C[1]),
            'FWHMG': Parameter(value=fwhmg, min=0.01),
            'FWHML': Parameter(value=fwhml, min=0.01),
            'scale': Parameter(value=scale, min=0, vary=racah)
        }
        if N is not None:
            pars['N'] = Parameter(value=N, vary=False)
            pars['Offset'] = Parameter(value=offset)
            pars['Poisson'] = Parameter(value=poisson, min=0, max=1)
            self.f = self.fShifted
        else:
            self.f = self.fUnshifted
        pars = {**pars, **self.intensities}

        self.params = pars

        if I < 1.5 or J1 < 1.5:
            self.params['Cl'].vary = False
        if I < 1.5 or J2 < 1.5:
            self.params['Cu'].vary = False
        if I < 1 or J1 < 1:
            self.params['Bl'].vary = False
        if I < 1 or J2 < 1:
            self.params['Bu'].vary = False
        if I == 0 or J1 == 0:
            self.params['Al'].vary = False
        if I == 0 or J2 == 0:
            self.params['Au'].vary = False

    def fUnshifted(self, x: ArrayLike) -> ArrayLike:
        """Calculate the response for an unshifted spectrum

        Parameters
        ----------
        x : ArrayLike

        Returns
        -------
        ArrayLike
        """
        centroid = self.params['centroid'].value
        Al = self.params['Al'].value
        Au = self.params['Au'].value
        Bl = self.params['Bl'].value
        Bu = self.params['Bu'].value
        Cl = self.params['Cl'].value
        Cu = self.params['Cu'].value
        FWHMG = self.params['FWHMG'].value
        FWHML = self.params['FWHML'].value
        scale = self.params['scale'].value

        result = np.zeros(len(x))
        x = self.transform(x)
        for line in self.lines:
            pos = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                    line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                        line]
            result += scale * self.params['Amp' + line].value * self.peak(
                x - pos, FWHMG, FWHML)

        return result

    def fShifted(self, x: ArrayLike) -> ArrayLike:
        """Calculate the response with :attr:`N` sidepeaks with an offset
        of :attr:`offset`

        Parameters
        ----------
        x : ArrayLike

        Returns
        -------
        ArrayLike
        """
        centroid = self.params['centroid'].value
        Al = self.params['Al'].value
        Au = self.params['Au'].value
        Bl = self.params['Bl'].value
        Bu = self.params['Bu'].value
        Cl = self.params['Cl'].value
        Cu = self.params['Cu'].value
        FWHMG = self.params['FWHMG'].value
        FWHML = self.params['FWHML'].value
        scale = self.params['scale'].value
        N = self.params['N'].value
        offset = self.params['Offset'].value
        poisson = self.params['Poisson'].value

        result = np.zeros(len(x))
        x = self.transform(x)
        for line in self.lines:
            pos = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                    line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                        line]
            for i in range(N + 1):
                result += self.params['Amp' + line].value * self.peak(
                    self.transform(x - i * offset) - pos, FWHMG,
                    FWHML) * (poisson**i) / np.math.factorial(i)
            result *= scale

        return result

    def peak(self, x: ArrayLike, FWHMG: float, FWHML: float) -> ArrayLike:
        """Calculates the Voigt profile given the Gaussian
        and Lorentzian FWHM

        Parameters
        ----------
        x : ArrayLike
            Evaluation points
        FWHMG : float
            Gaussian FWHM
        FWHML : float
            Lorentzian FWHM

        Returns
        -------
        ArrayLike
        """
        sigma, gamma = FWHMG / sqrt2log2t2, FWHML / 2
        return voigt_profile(x, sigma, gamma) / voigt_profile(0, sigma, gamma)

    def calcShift(self, I: float, J: float, F: int) -> ArrayLike:
        """Calculate the coefficients for the energy shift due to the hyperfine
        interaction up to the octupole moment. A general equation is used
        so extending to higher orders is possible.

        Parameters
        ----------
        I : float
            Nuclear spin
        J : float
            Electronic spin
        F : int
            Hyperfine level spin

        Returns
        -------
        ArrayLike
            Individual coefficients, in ascending order
        """
        phase = (-1)**(I + J + F)
        contrib = []
        for k in range(1, 4):
            n = float(wigner_6j(I, J, F, J, I, k))
            d = float(
                wigner_3j(I, k, I, -I, 0, I) * wigner_3j(J, k, J, -J, 0, J))
            shift = phase * n / d
            if not np.isfinite(shift):
                contrib.append(0)
            else:
                if k == 1:
                    shift = shift * (I * J)
                elif k == 2:
                    shift = shift / 4
                contrib.append(shift)
        return contrib

    def pos(self) -> ArrayLike:
        """Returns the positions of the peaks in MHz in the hyperfine spectrum

        Returns
        -------
        ArrayLike
        """
        centroid = self.params['centroid'].value
        Al = self.params['Al'].value
        Au = self.params['Au'].value
        Bl = self.params['Bl'].value
        Bu = self.params['Bu'].value
        Cl = self.params['Cl'].value
        Cu = self.params['Cu'].value
        pos = []
        for line in self.lines:
            p = centroid + Au * self.scaling_Au[line] + Bu * self.scaling_Bu[
                line] + Cu * self.scaling_Cu[line] - Al * self.scaling_Al[
                    line] - Bl * self.scaling_Bl[line] - Cl * self.scaling_Cl[
                        line]
            pos.append(p)
        return pos

    def calculateFWHM(self) -> Tuple[float, float]:
        """Calculate the total FWHM of the profiles, with uncertainty,
        taking the correlations into account.

        Returns
        -------
        Tuple[float, float]
            Tuple of the form (value, uncertainty)
        """
        G, Gu = self.params['FWHMG'].value, self.params['FWHMG'].unc
        L, Lu = self.params['FWHML'].value, self.params['FWHML'].unc
        try:
            correl = self.params['FWHMG'].correl['FWHML']
        except KeyError:
            correl = 0
        G, L = unc.correlated_values_norm([(G, Gu), (L, Lu)],
                                          np.array([[1, correl], [correl, 1]]))
        fwhm = 0.5346 * L + (0.2166 * L * L + G * G)**0.5
        return fwhm.nominal_value, fwhm.std_dev
