import numpy as np
try:
    from scipy.integrate import simpson
except ImportError:
    from scipy.integrate import simps as simpson

from ..math.orientations import zcw_spherical as crystallites


__all__ = ['CSA', 'Pake', 'SecCentralLine', 'CzjzekCT']


def _make_broadening(x: np.ndarray, sigma: float, mode: str):
    dx = x[1] - x[0]
    _x = np.arange(len(x)) * dx
    _x -= 0.5 * _x[-1]
    if mode == 'l':
        apd = 2 * sigma / (4*_x**2 + sigma**2) / np.pi
    else:
        ln2 = np.log(2)
        apd = np.exp(-4*ln2 * (_x/sigma)**2) * 2 * np.sqrt(ln2/np.pi) / sigma
    return apd


def _make_bins(x: np.ndarray) -> np.ndarray:
    bins = 0.5 * (x[1:] + x[:-1])
    return np.r_[0.5 * (-x[1] + 3 * x[0]), bins, 0.5 * (3 * x[-1] - x[-2])]


def _make_x(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    _x = x
    dx = x[1:] - x[:-1]
    dx = np.min(dx)
    width = x[-1] - x[0]
    _x = np.arange(width/dx - 1) * dx + x[0]

    bins = (_x[1:] + _x[:-1]) / 2
    bins = np.r_[_x[0]-dx/2, bins, _x[-1] + dx/2]
    return _x, bins


def _make_quad_prefactor(cq, f_l, spin):
    omega_q = 2 * np.pi * cq / (2 * spin * (2 * spin - 1))
    return 1.5 * (omega_q ** 2 / (2 * np.pi * f_l)) * (spin * (spin + 1) - 0.75)


class Pake:
    type = 'Spectrum'
    name = 'Pake'
    equation = ''
    params = ['A', r'\delta', r'\eta', r'\Sigma_{B}', r't_{pulse}']
    bounds = [(0, None), (0, None), (0, 1), (0, None)]
    choices = [('Broadening', 'broad', {'Gaussian': 'g', 'Lorentzian': 'l'})]

    @staticmethod
    def func(
            x: np.ndarray,
            c: float,
            delta: float,
            eta: float,
            sigma: float,
            t_pulse: float,
            broad: str = 'g',
    ) -> np.ndarray:
        a, b, _ = crystallites(100000)
        omega = delta * 0.5 * (3*np.cos(b)**2 - 1 - eta * np.sin(b)**2 * np.cos(2*a))
        x_used, bins = _make_x(x)

        s_left = np.histogram(omega, bins=bins)[0]
        s_right = np.histogram(-omega, bins=bins)[0]
        s = s_left + s_right

        if sigma != 0:
            apd = _make_broadening(x_used, sigma, broad)
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        omega_1 = np.pi/2/t_pulse
        attn = omega_1 * np.sin(t_pulse*np.sqrt(omega_1**2 + 0.5*(2*np.pi*x_used)**2)) / np.sqrt(omega_1**2 + (np.pi*x_used)**2)

        ret_val *= attn
        ret_val /= simpson(y=ret_val, x=x_used)

        if x_used.size == x.size:
            return c * ret_val
        else:
            return c * np.interp(x=x, xp=x_used, fp=ret_val)


class CSA:
    type = 'Spectrum'
    name = 'CSA'
    equation = ''
    params = ['A', r'\delta', r'\eta', r'\omega_{iso}', r'\Sigma_{B}']
    bounds = [(0, None), (0, None), (0, 1), (None, None), (0, None)]
    choices = [('Broadening', 'broad', {'Gaussian': 'g', 'Lorentzian': 'l'})]

    @staticmethod
    def func(
            x: np.ndarray,
            c: float,
            delta: float,
            eta: float,
            w_iso: float,
            sigma: float,
            broad: str = 'g',
    ) -> np.ndarray:
        a, b, _ = crystallites(100000)

        omega = w_iso + delta * 0.5 * (3*np.cos(b)**2 - 1 - eta * np.sin(b)**2 * np.cos(2*a))

        s = np.histogram(omega, bins=_make_bins(x))[0]

        if sigma != 0:
            print(len(s))
            apd = _make_broadening(x, sigma, broad)
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        return c * ret_val / simpson(y=ret_val, x=x)


class SecCentralLine:
    type = 'Spectrum'
    name = 'Central Transition (2nd order)'
    equation = ''
    params = ['A', 'C_{Q}', r'\eta', r'\omega_{iso}', 'GB', r'\omega_{L}']
    bounds = [(0, None), (0, None), (0, 1), (None, None), (0, None), (0, None)]
    choices = [('Spin', 'spin', {'3/2': 1.5, '5/2': 2.5}),
               ('Broadening', 'broad', {'Gaussian': 'g', 'Lorentzian': 'l'})]

    @staticmethod
    def func(
            x: np.ndarray,
            c: float,
            cq: float,
            eta: float,
            f_iso: float,
            gb: float,
            f_l: float,
            spin: float = 2.5,
            broad: str = 'g',
    ) -> np.ndarray:
        a, b, _ = crystallites(200000)

        # coupling constant
        coupling = _make_quad_prefactor(cq, f_l, spin)

        # orientation
        cos2phi = np.cos(2*a)
        cos_theta_square = 0.5 + 0.5 * np.cos(2*b)  # cos^2(x) = 1/2 (1+cos(2x)
        prefactor_a = -3.375 + 2.25 * eta * cos2phi - 0.375 * (eta*cos2phi)**2
        prefactor_b = 3.75 - 0.5 * eta**2 - 2 * eta*cos2phi + 0.75 * (eta*cos2phi)**2
        prefactor_c = -0.375 + (eta**2) / 3 - 0.25 * eta*cos2phi - 0.375 * (eta*cos2phi)**2

        orient = np.zeros_like(cos2phi)
        orient += prefactor_a * cos_theta_square ** 2
        orient += prefactor_b * cos_theta_square
        orient += prefactor_c

        omega = 2*np.pi*f_iso + coupling * orient
        s = np.histogram(omega / (2*np.pi), bins=_make_bins(x))[0]

        if gb != 0:
            apd = _make_broadening(x, gb, broad)
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        return c * ret_val / simpson(y=ret_val, x=x)


class CzjzekCT:
    type = 'Spectrum'
    name = 'Czjzek (Central Line)'
    equation = ''
    params = ['A', r'\sigma', r'f_{iso}', 'GB', r'\omega_{L}']
    bounds = [(0, None), (0, None), (None, None), (0, None), (0, None)]
    choices = [('Spin', 'spin', {'3/2': 1.5, '5/2': 2.5}),
               ('Broadening', 'broad', {'Gaussian': 'g', 'Lorentzian': 'l'})]

    @staticmethod
    def func(
            x: np.ndarray,
            c: float,
            sigma: float,
            f_iso: float,
            gb: float,
            f_l: float,
            spin: float = 1.5,
            broad: str = 'g',
    ) -> np.ndarray:

        cq = np.linspace(0, 5 * sigma, num=200)
        eta = np.linspace(0, 1, num=200)

        a, b, _ = crystallites(10000)
        a = a.reshape(-1, 1)
        b = b.reshape(-1, 1)

        # coupling constant
        dist = CzjzekCT.czjzek(cq, eta[:, None], sigma)

        coupling = _make_quad_prefactor(cq, f_l, spin)

        # orientation
        e_times_phi = eta * np.cos(2 * a)
        e_times_phi_sq = e_times_phi * e_times_phi

        cos_theta_square = 0.5 + 0.5 * np.cos(2 * b)  # cos^2(x) = 1/2 (1+cos(2x)
        prefactor_a = -3.375 + 2.25 * e_times_phi - 0.375 * e_times_phi_sq
        prefactor_b = 3.75 - 0.5 * eta ** 2 - 2 * e_times_phi + 0.75 * e_times_phi_sq
        prefactor_c = -0.375 + (eta ** 2) / 3 - 0.25 * e_times_phi - 0.375 * e_times_phi_sq

        orient = np.zeros((a.size, eta.size))
        orient += prefactor_a * cos_theta_square ** 2
        orient += prefactor_b * cos_theta_square
        orient += prefactor_c

        vQ = -(orient[..., None] * coupling[None, None, :]) + f_iso
        weights = np.ones_like(vQ) * dist

        bins = _make_bins(x)

        s, _ = np.histogram(
            vQ.reshape(-1),
            weights=weights.reshape(-1),
            bins=bins,
        )

        if gb != 0:
            apd = _make_broadening(x, gb, broad)
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        return c * ret_val / simpson(y=ret_val, x=x)

    @staticmethod
    def czjzek(cq, eta, sigma):
        return np.exp(-cq ** 2 * (1 + eta ** 2 / 3) / 2 / sigma ** 2) * cq ** 4 * eta * (
                    1 - eta ** 2 / 9) * np.sqrt(2 / np.pi) / sigma ** 5
