import numpy as np
from astropy import units as u, constants as ac

from Helpers import *

DY_POLARIZABILITY = 184.4 # in a.u, most precise measured value of Dy polarizability
DY_MASS = 164*u.u
DY_DIPOLE_MOMENT = 9.93 * ac.muB

#####################################################################
# POTENTIALS                                                        #
#####################################################################

def gravitational_potential(positions, m):
    return m * ac.g0 * positions

def single_gaussian_beam_potential(positions, waists, alpha = DY_POLARIZABILITY, P=1, wavelength=1.064*u.um):
    A = 2*P/(np.pi*w(positions[1,:], waists[0], wavelength)*w(positions[1,:], waists[1], wavelength))
    U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U = - U_tilde * A * np.exp(-2 * ((positions[0,:]/w(positions[1,:], waists[0], wavelength))**2 + (positions[2,:]/w(positions[1,:], waists[1], wavelength))**2)) 
    return U

def astigmatic_single_gaussian_beam_potential(positions, waists, del_y, alpha = DY_POLARIZABILITY, P=1, wavelength=1.064*u.um):
    A = 2*P/(np.pi*w(positions[1,:] - (del_y/2), waists[0], wavelength)*w(positions[1,:] + (del_y/2), waists[1], wavelength))
    U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U = - U_tilde * A * np.exp(-2 * ((positions[0,:]/w(positions[1,:] - (del_y/2), waists[0], wavelength))**2 + (positions[2,:]/w(positions[1,:] + (del_y/2), waists[1], wavelength))**2)) 
    return U
    
def modulated_single_gaussian_beam_potential(positions, waists, alpha = DY_POLARIZABILITY, P=1, wavelength=1.064*u.um, mod_amp=1):
    mod_amp = mod_amp * waists[0]
    n_points = len(positions[0,:])
    dx, x_mod = modulation_function(mod_amp, n_points, func = 'arccos')
    A = 2*P/(np.pi*w(positions[1,:], waists[0], wavelength)*w(positions[1,:], waists[1], wavelength))
    U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    dU = np.zeros(2*n_points)
    for i in range(len(x_mod)):
        dU = np.vstack((dU, np.exp(-2 * (np.subtract(x_mod[i], positions[0,:])/w(positions[1,:], waists[0], wavelength))**2)))
    U = - U_tilde * A * 1/(2*mod_amp) * np.trapz(dU, dx = dx, axis = 0)
    return U

def harmonic_potential(pos, v, xoffset, yoffset, m = DY_MASS):
    U_Harmonic = ((0.5 * m * (2 * np.pi * v*u.Hz)**2 * (pos*u.um - xoffset*u.um)**2)/ac.k_B).to(u.uK) + yoffset*u.uK
    return U_Harmonic.value

def gaussian_potential(pos, amp, waist, xoffset, yoffset):
    U_Gaussian = amp * np.exp(-2 * ((pos + xoffset) / waist)**2) + yoffset
    return U_Gaussian

def crossed_beam_potential(positions, waists, P, options, alpha = DY_POLARIZABILITY, wavelength=1.064*u.um):
    
    delta = options['delta']

    foci_shift = options['foci_shift']
    focus_shift_beam_1 = foci_shift[0]
    focus_shift_beam_2 = foci_shift[1]

    beam_disp = options['beam_disp']
    beam_1_disp = (np.ones(np.shape(positions.T)) * np.array(beam_disp[0])).T * beam_disp[0].unit
    beam_2_disp = (np.ones(np.shape(positions.T)) * np.array(beam_disp[1])).T * beam_disp[1].unit

    beam_1_positions = positions + beam_1_disp
    A_1 = 2*P[0]/(np.pi*w(beam_1_positions[1,:] + focus_shift_beam_1, waists[0][0], wavelength)*w(beam_1_positions[1,:] + focus_shift_beam_1, waists[0][1], wavelength))
    U_1_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U_1 = - U_1_tilde * A_1 * np.exp(-2 * ((beam_1_positions[0,:]/w(beam_1_positions[1,:] + focus_shift_beam_1, waists[0][0], wavelength))**2 + (beam_1_positions[2,:]/w(beam_1_positions[1,:] + focus_shift_beam_1, waists[0][1], wavelength))**2)) 
    
    R = rotation_matrix([0, 0, 1], np.radians(delta))
    beam_2_positions = np.dot(R, positions + beam_2_disp)
    A_2 = 2*P[1]/(np.pi*w(beam_2_positions[1,:] + focus_shift_beam_2, waists[1][0], wavelength)*w(beam_2_positions[1,:] + focus_shift_beam_2, waists[1][1], wavelength))
    U_2_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U_2 = - U_2_tilde * A_2 * np.exp(-2 * ((beam_2_positions[0,:]/w(beam_2_positions[1,:] + focus_shift_beam_2, waists[1][0], wavelength))**2 + (beam_2_positions[2,:]/w(beam_2_positions[1,:] + focus_shift_beam_2, waists[1][1], wavelength))**2)) 

    U = U_1 + U_2

    return U

def astigmatic_crossed_beam_potential(positions, waists, P, options, alpha = DY_POLARIZABILITY, wavelength=1.064*u.um):

    delta = options['delta']

    del_y = options['foci_disp_crossed']
    del_y_1 = del_y[0]
    del_y_2 = del_y[1]

    
    foci_shift = options['foci_shift']
    focus_shift_beam_1 = foci_shift[0]
    focus_shift_beam_2 = foci_shift[1]

    beam_disp = options['beam_disp']
    beam_1_disp = (np.ones(np.shape(positions.T)) * np.array(beam_disp[0])).T * beam_disp[0].unit
    beam_2_disp = (np.ones(np.shape(positions.T)) * np.array(beam_disp[1])).T * beam_disp[1].unit

    beam_1_positions = positions + beam_1_disp
    A_1 = 2*P[0]/(np.pi*w(beam_1_positions[1,:] - (del_y_1/2) + focus_shift_beam_1, waists[0][0], wavelength)*w(beam_1_positions[1,:] + (del_y_1/2) + focus_shift_beam_1, waists[0][1], wavelength))
    U_1_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U_1 = - U_1_tilde * A_1 * np.exp(-2 * ((beam_1_positions[0,:]/w(beam_1_positions[1,:] - (del_y_1/2) + focus_shift_beam_1, waists[0][0], wavelength))**2 + (beam_1_positions[2,:]/w(beam_1_positions[1,:] + (del_y_1/2) + focus_shift_beam_1, waists[0][1], wavelength))**2)) 
    
    R = rotation_matrix([0, 0, 1], np.radians(delta))
    beam_2_positions = np.dot(R, positions + beam_2_disp)
    A_2 = 2*P[1]/(np.pi*w(beam_2_positions[1,:] - (del_y_2/2) + focus_shift_beam_2, waists[1][0], wavelength)*w(beam_2_positions[1,:] + (del_y_2/2) + focus_shift_beam_2, waists[1][1], wavelength))
    U_2_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
    U_2 = - U_2_tilde * A_2 * np.exp(-2 * ((beam_2_positions[0,:]/w(beam_2_positions[1,:] - (del_y_2/2) + focus_shift_beam_2, waists[1][0], wavelength))**2 + (beam_2_positions[2,:]/w(beam_2_positions[1,:] + (del_y_2/2) + focus_shift_beam_2, waists[1][1], wavelength))**2)) 

    U = U_1 + U_2

    return U

