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

from ..math.orientations import zcw_spherical as crystallites


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, c, delta, eta, sigma, t_pulse, broad='g'):
        a, b, _ = crystallites(100000)
        bins = 0.5 * (x[1:] + x[:-1])
        bins = np.r_[0.5*(3*x[0]-x[1]), bins, 0.5*(3*x[-1]-x[-2])]

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

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

        if sigma != 0:
            _x = np.arange(len(x))*(x[1]-x[0])
            _x -= 0.5*_x[-1]

            if broad == 'l':
                apd = 2 * sigma / (4 * _x**2 + sigma**2) / pi
            else:
                apd = np.exp(-4 * np.log(2) * (_x/sigma)**2) * 2 * np.sqrt(np.log(2) / pi) / sigma

            ret_val = np.convolve(s, apd, mode='same')

        else:
            ret_val = s
            
        omega_1 = pi/2/t_pulse
        attn = omega_1 * np.sin(t_pulse*np.sqrt(omega_1**2+0.5*(2*pi*x)**2)) / \
            np.sqrt(omega_1**2+(np.pi*x)**2)

        ret_val *= attn

        return c * ret_val / simpson(ret_val, x)


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, c, delta, eta, w_iso, sigma, broad='g'):
        a, b, _ = crystallites(100000)
        bins = 0.5 * (x[1:] + x[:-1])
        bins = np.r_[0.5*(-x[1] + 3*x[0]), bins, 0.5*(3*x[-1] - x[-2])]

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

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

        if sigma != 0:
            _x = np.arange(len(x)) * (x[1] - x[0])
            _x -= 0.5 * _x[-1]
            if broad == 'l':
                apd = 2 * sigma / (4*_x**2 + sigma**2) / pi
            else:
                apd = np.exp(-4 * np.log(2) * (_x / sigma) ** 2) * 2 * np.sqrt(np.log(2) / pi) / sigma
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        return c * ret_val / simpson(ret_val, 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, c, cq, eta, f_iso, gb, f_l, spin=2.5, broad='g'):
        a, b, _ = crystallites(200000)
        bins = 0.5 * (x[1:] + x[:-1])
        bins = np.r_[0.5*(-x[1] + 3*x[0]), bins, 0.5*(3*x[-1] - x[-2])]

        # coupling constant
        omega_q = 2 * np.pi * cq / (2*spin*(2*spin-1))
        coupling = 1.5 * (omega_q**2 / (2*np.pi*f_l)) * (spin*(spin+1)-0.75)

        # 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=bins)[0]

        if gb != 0:
            _x = np.arange(len(x)) * (x[1]-x[0])
            _x -= 0.5*_x[-1]
            if broad == 'l':
                apd = 2*gb / (4*_x**2 + gb**2) / np.pi
            else:
                apd = np.exp(-4*np.log(2) * (_x/gb)**2) * 2 * np.sqrt(np.log(2)/np.pi) / gb
            ret_val = np.convolve(s, apd, mode='same')
        else:
            ret_val = s

        return c * ret_val / simpson(ret_val, x)
