# -*- coding: utf-8 -*-
"""
qspec.models._spectrum
======================

Spectrum classes for lineshape models.
"""

import numpy as np
from scipy.stats import norm, cauchy
from scipy.special import voigt_profile
from scipy.special import wofz

from qspec.physics import source_energy_pdf
from qspec.models._base import Model

__all__ = ['SPECTRA', 'fwhm_voigt', 'fwhm_voigt_d', 'Spectrum', 'Gauss', 'Lorentz', 'LorentzQI',
           'Voigt', 'VoigtDerivative', 'VoigtAsy', 'VoigtCEC', 'GaussChi2']


# The names of the spectra. Includes all spectra that appear in the GUI.
SPECTRA = ['Gauss', 'Lorentz', 'Voigt', 'VoigtDerivative', 'VoigtAsy', 'VoigtCEC', 'GaussChi2']


def fwhm_voigt(gamma, sigma):
    """
    :param gamma: The FWHM of the Lorentzian.
    :param sigma: The standard deviation of the Gaussian.
    :returns: The FWHM of a Voigt profile.
    """
    f_l = abs(gamma)
    f_g = abs(np.sqrt(8 * np.log(2)) * sigma)
    return 0.5346 * f_l + np.sqrt(0.2166 * f_l ** 2 + f_g ** 2)


def fwhm_voigt_d(gamma, gamma_d, sigma, sigma_d):
    """
    :param gamma: The FWHM of the Lorentzian.
    :param gamma_d: The uncertainty of 'gamma'.
    :param sigma: The standard deviation of the Gaussian.
    :param sigma_d: The uncertainty of 'sigma'.
    :returns: The uncertainty of the FWHM of a Voigt profile.
    """
    f_l = abs(gamma)
    f_l_d = abs(gamma_d)
    f_g = abs(np.sqrt(8 * np.log(2)) * sigma)
    f_g_d = abs(np.sqrt(8 * np.log(2)) * sigma_d)
    f_g_d = f_g / np.sqrt(0.2166 * f_l ** 2 + f_g ** 2) * f_g_d
    f_l_d = (0.5346 + 0.2166 * f_l / np.sqrt(0.2166 * f_l ** 2 + f_g ** 2)) * f_l_d
    return np.sqrt(f_l_d ** 2 + f_g_d ** 2)


class Spectrum(Model):
    def __init__(self):
        super().__init__(model=None)
        self.type = 'Spectrum'

    def evaluate(self, x, *args, **kwargs):
        return np.zeros_like(x)

    @property
    def dx(self):
        return self.fwhm() * 0.02

    def fwhm(self):
        """
        :returns: The full width at half maximum (FWHM) of a spectrum.
        """
        return 1.

    def min(self):
        return -5 * self.fwhm()

    def max(self):
        return 5 * self.fwhm()


class Lorentz(Spectrum):
    def __init__(self):
        """
        A Lorentzian peak function.
        """
        super().__init__()
        self.type = 'Lorentz'

        self._add_arg('Gamma', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        scale = 0.5 * args[0]
        return np.pi * scale * cauchy.pdf(x, loc=0, scale=scale)

    def fwhm(self):
        return abs(self.vals[self.p['Gamma']])

    def min(self):
        return -10 * self.fwhm()

    def max(self):
        return 10 * self.fwhm()


class LorentzQI(Spectrum):
    def __init__(self):
        """
        A Lorentzian peak function including a dispersive term for consideration of quantum interference (QI) effects.
        """
        super().__init__()
        self.type = 'LorentzQI'

        self._add_arg('Gamma', 1., False, False)

    def evaluate_qi(self, x, x_qi, *args, **kwargs):
        """
        :param x: The input values.
        :param x_qi: The x value of the interfering Lorentzian.
        :param args: The function parameters. Must have length self.size.
        :param kwargs: Additional keyword arguments.
        :returns: The function results of the dispersive Lorentzian at the input values 'x'.
        """
        scale = 0.5 * args[0]
        return 2 * scale ** 2 * np.real(1 / ((x + 1j * scale) * (x - x_qi - 1j * scale)))

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        scale = 0.5 * args[0]
        return scale * np.pi * cauchy.pdf(x, loc=0, scale=scale)

    def fwhm(self):
        return abs(self.vals[self.p['Gamma']])

    def min(self):
        return -10 * self.fwhm()

    def max(self):
        return 10 * self.fwhm()


class Gauss(Spectrum):
    def __init__(self):
        """
        A Gaussian peak function.
        """
        super().__init__()
        self.type = 'Gauss'

        self._add_arg('sigma', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        return np.sqrt(2 * np.pi) * args[0] * norm.pdf(x, loc=0, scale=args[0])

    def fwhm(self):
        return abs(np.sqrt(8 * np.log(2)) * self.vals[self.p['sigma']])

    def min(self):
        return -2.5 * self.fwhm()

    def max(self):
        return 2.5 * self.fwhm()


class Voigt(Spectrum):
    def __init__(self):
        """
        A Voigt peak function.
        """
        super().__init__()
        self.type = 'Voigt'

        self._add_arg('Gamma', 1., False, False)
        self._add_arg('sigma', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        return voigt_profile(x, args[1], 0.5 * args[0]) / voigt_profile(0, args[1], 0.5 * args[0])

    def fwhm(self):
        return fwhm_voigt(self.vals[self.p['Gamma']], self.vals[self.p['sigma']])


class VoigtDerivative(Spectrum):
    def __init__(self):
        """
        The first derivative of a Voigt peak function.
        """
        super().__init__()
        self.type = 'VoigtDerivative'

        self._add_arg('Gamma', 1., False, False)
        self._add_arg('sigma', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum. (Pos of min/max to be determined).
        z = (x + 1j * 0.5 * args[0]) / (np.sqrt(2) * args[1])
        fwhm = abs(0.5346 * args[0] + np.sqrt(0.2166 * args[0] ** 2 + args[1] ** 2))  # for normalization
        return -(z * wofz(z)).real / (np.sqrt(np.pi) * args[1] ** 2) / voigt_profile(0, args[1], 0.5 * args[0]) * fwhm

    def fwhm(self):
        return fwhm_voigt(self.vals[self.p['Gamma']], self.vals[self.p['sigma']])


class VoigtAsy(Spectrum):
    def __init__(self):
        """
        An asymmetric Voigt peak function.
        """
        super().__init__()
        self.type = 'VoigtAsy'

        self._add_arg('Gamma', 1., False, False)
        self._add_arg('sigma', 1., False, False)
        self._add_arg('asyPar', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        gamma = args[0] / (1 + np.exp(args[2] * x))
        return voigt_profile(x, args[1], gamma) / voigt_profile(0, args[1], 0.5 * args[0])

    def fwhm(self):
        return fwhm_voigt(self.vals[self.p['Gamma']], self.vals[self.p['sigma']])


class VoigtCEC(Spectrum):
    def __init__(self):
        """
        A series of `n` Voigt peak functions shifted by `i * shift` and scaled by `i * ratio`, where `0 <= i <= n`.
        """
        super().__init__()
        self.type = 'VoigtCEC'

        self._add_arg('Gamma', 1., False, False)
        self._add_arg('sigma', 1., False, False)
        self._add_arg('shift', 0., False, False)
        self._add_arg('ratio', 0., False, False)
        self._add_arg('n', 0, True, False)

    def evaluate(self, x, *args, **kwargs):  # Normalize to the maximum.
        return np.sum([args[3] ** i * voigt_profile(x - i * args[2], args[1], 0.5 * args[0])
                       for i in range(int(args[4]) + 1)], axis=0) / voigt_profile(0, args[1], 0.5 * args[0])

    def fwhm(self):
        return fwhm_voigt(self.vals[self.p['Gamma']], self.vals[self.p['sigma']])


def _gauss_chi2_taylor_fwhm(sigma, xi):
    """
    :param sigma: The standard deviation of the Gaussian.
    :param xi: The asymmetry parameter.
    :returns: An estimated Taylor expansion of the FWHM of the 'GaussChi2' model.
    """
    a = [2.70991004e+00, 2.31314470e-01, -8.11610976e-02, -1.48897229e-02, 1.11677618e-01,  # order 1, 2
         8.06412708e-03, -6.59788156e-04, -1.06067275e-02, 1.47400503e-03]  # order 3
    return a[0] * sigma + a[1] * xi + (a[2] * sigma ** 2 + a[3] * xi ** 2 + a[4] * sigma * xi) / 2 \
        + (a[5] * sigma ** 3 + a[6] * xi ** 3 + a[7] * sigma ** 2 * xi + a[8] * sigma * xi ** 2) / 6


def _gauss_chi2_fwhm(sigma, xi):
    """
    :param sigma: The standard deviation of the Gaussian.
    :param xi: The asymmetry parameter.
    :returns: An estimation of the FWHM of the 'GaussChi2' model.
    """
    x = (sigma, xi)
    a, b, c = (1.39048239, 0.49098627, 0.61375536)
    return a * x[0] * (np.arctan(b * (np.abs(x[1] / x[0]) ** c - 1)) + np.arctan(b)) + np.sqrt(8 * np.log(2)) * x[0]


class GaussChi2(Spectrum):  # TODO: The fwhm of GaussChi2 is fitting quite well now, but could be improved further.
    """
    An analytical convolution of a Gaussian with a Boltzmann distribution.
    This could be the velocity distribution of a thermalized ion ensemble produced
    on a linear or noisy voltage potential.
    """
    def __init__(self):
        super().__init__()
        self.type = 'GaussChi2'

        self._add_arg('sigma', 1., False, False)
        self._add_arg('xi', 1., False, False)

    def evaluate(self, x, *args, **kwargs):  # Maximum unknown, normalize to 0 position, f(0) ~> 0.8 * max(f(x)).
        return source_energy_pdf(x, 0, args[0], args[1], collinear=True) \
            / source_energy_pdf(0, 0, args[0], args[1], collinear=True)

    def fwhm(self):
        sigma = np.abs(self.vals[self.p['sigma']])
        xi = np.abs(self.vals[self.p['xi']])
        if xi == 0:
            return np.sqrt(8 * np.log(2)) * sigma
        return _gauss_chi2_fwhm(sigma, xi)
        # return _gauss_chi2_taylor_fwhm(sigma, xi)

    def min(self):
        sigma = np.abs(self.vals[self.p['sigma']])
        xi = self.vals[self.p['xi']]
        if xi < 0:
            return -5 * np.sqrt(2 * np.log(2)) * sigma
        return -5 * (np.sqrt(2 * np.log(2)) * sigma + xi)

    def max(self):
        sigma = np.abs(self.vals[self.p['sigma']])
        xi = self.vals[self.p['xi']]
        if xi > 0:
            return 5 * np.sqrt(2 * np.log(2)) * sigma
        return 5 * (np.sqrt(2 * np.log(2)) * sigma - xi)
