import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from scipy import signal, interpolate
from scipy.optimize import curve_fit
from astropy import units as u, constants as ac

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

#####################################################################
# HELPER FUNCTIONS                                                  #
#####################################################################

def orderOfMagnitude(number):
    return math.floor(math.log(number, 10))

def rotation_matrix(axis, theta):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.
    In 2-D it is just,
        thetaInRadians = np.radians(theta)
        c, s = np.cos(thetaInRadians), np.sin(thetaInRadians)
        R = np.array(((c, -s), (s, c)))
    In 3-D, one way to do it is use the Euler-Rodrigues Formula as is done here   
    """
    axis = np.asarray(axis)
    axis = axis / math.sqrt(np.dot(axis, axis))
    a = math.cos(theta / 2.0)
    b, c, d = -axis * math.sin(theta / 2.0)
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
                     [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
                     [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx

def modulation_function(mod_amp, n_points, func = 'arccos'):
    
    if func == 'sin':
        phi         = np.linspace(0, 2*np.pi, n_points)
        mod_func    = mod_amp * np.sin(phi)
    elif func == 'arccos':
        # phi        = np.linspace(0, 2*np.pi, n_points)
        # mod_func   = mod_amp * (2/np.pi * np.arccos(phi/np.pi-1) - 1)
        phi      = np.linspace(0, 2*np.pi, int(n_points/2))
        tmp_1    = 2/np.pi * np.arccos(phi/np.pi-1) - 1
        tmp_2    = np.flip(tmp_1)
        mod_func = mod_amp * np.concatenate((tmp_1, tmp_2))
    elif func == 'triangle':
        phi         = np.linspace(0, 2*np.pi, n_points)
        mod_func    = mod_amp * signal.sawtooth(phi, width = 0.5) # width of 0.5 gives symmetric rising triangle ramp
    elif func == 'square':
        phi         = np.linspace(0, 1.99*np.pi, n_points)
        mod_func    = mod_amp * signal.square(phi, duty = 0.5)
    else:
        mod_func = None    
    
    if mod_func is not None:
        dx          = (max(mod_func) - min(mod_func))/(2*n_points)

    return dx, mod_func

#####################################################################
# BEAM PARAMETERS                                                   #
#####################################################################

# Rayleigh length
def z_R(w_0, lamb)->np.ndarray:
    return np.pi*w_0**2/lamb

# Beam Radius
def w(pos, w_0, lamb):
    return w_0*np.sqrt(1+(pos / z_R(w_0, lamb))**2)

#####################################################################
# COLLISION RATES, PSD                                              #
#####################################################################

def calculateAtomNumber(NCount, pixel_size, magnification, eta):
    sigma = 8.468e-14 * (u.m)**(2)
    return (1/eta * 1/sigma * NCount * pixel_size**2/magnification**2).decompose()

def meanThermalVelocity(T, m = DY_MASS):
    return 4 * np.sqrt((ac.k_B * T) /(np.pi * m))

def particleDensity(w_x, w_z, Power, N, T, m = DY_MASS): # For a thermal cloud
    v_x = calculateTrapFrequency(w_x, w_z, Power, dir = 'x')
    v_y = calculateTrapFrequency(w_x, w_z, Power, dir = 'y')
    v_z = calculateTrapFrequency(w_x, w_z, Power, dir = 'z')
    return N * (2 * np.pi)**3 * (v_x * v_y * v_z) *  (m / (2 * np.pi * ac.k_B * T))**(3/2)

def calculateParticleDensityFromMeasurements(v_x, dv_x, v_y, dv_y, v_z, dv_z, w_x, w_z, T_x, T_y, dT_x, dT_y, modulation_depth, N, m = DY_MASS):
        alpha_x =  [(v_x[0]/x)**(2/3) for x in v_x]
        dalpha_x = [alpha_x[i] * np.sqrt((dv_x[0]/v_x[0])**2 + (dv_x[i]/v_x[i])**2) for i in range(len(v_x))]
        alpha_y = [(v_z[0]/y)**2 for y in v_z]
        dalpha_y = [alpha_y[i] * np.sqrt((dv_z[0]/v_z[0])**2 + (dv_z[i]/v_z[i])**2) for i in range(len(v_z))]

        avg_alpha = [(g + h) / 2 for g, h in zip(alpha_x, alpha_y)]
        new_aspect_ratio = (w_x * avg_alpha) / w_z

        avg_T = [(g + h) / 2 for g, h in zip(T_x, T_y)]
        avg_dT = [0.5 * np.sqrt(g**2 + h**2) for g, h in zip(dT_x, dT_y)]

        pd = np.zeros(len(modulation_depth))
        dpd = np.zeros(len(modulation_depth))
        
        for i in range(len(modulation_depth)):
            particle_density = (N * (2 * np.pi)**3 * (v_x[i] * v_y[i] * v_z[i]) *  (m / (2 * np.pi * ac.k_B * avg_T[i]))**(3/2)).decompose()
            pd[i] = particle_density.value
            dpd[i] = (((N * (2 * np.pi)**3 * (m / (2 * np.pi * ac.k_B * avg_T[i]))**(3/2)) * ((dv_x[i] * v_y[i] * v_z[i]) + (v_x[i] * dv_y[i] * v_z[i]) + (v_x[i] * v_y[i] * dv_z[i]) - (1.5*(v_x[i] * v_y[i] * v_z[i])*(avg_dT[i]/avg_T[i])))).decompose()).value  
        
        pd = pd*particle_density.unit
        dpd = dpd*particle_density.unit
        
        return pd, dpd, avg_T, avg_dT, new_aspect_ratio
    
def thermaldeBroglieWavelength(T, m = DY_MASS):
    return np.sqrt((2*np.pi*ac.hbar**2)/(m*ac.k_B*T))

def scatteringLength(B, FR_choice = 1, ABKG_choice = 1):
    # Dy 164 a_s versus B in 0 to 8G range
    # should match SupMat of PhysRevX.9.021012, fig S5 and descrption
    # https://journals.aps.org/prx/supplemental/10.1103/PhysRevX.9.021012/Resubmission_Suppmat.pdf
    
    if FR_choice == 1: # new values
        
        if ABKG_choice == 1:
            a_bkg = 85.5 * ac.a0
        elif ABKG_choice == 2:
            a_bkg = 93.5 * ac.a0
        elif ABKG_choice == 3:
            a_bkg = 77.5 * ac.a0
        
        #FR resonances
        #[B11 B12 B2 B3 B4 B51 B52 B53 B6 B71 B72 B81 B82 B83 B9]
        resonanceB  = [1.295, 1.306, 2.174, 2.336, 2.591, 2.74, 2.803, 2.78, 3.357, 4.949, 5.083, 7.172, 7.204, 7.134, 76.9] * u.G #resonance position
        #[wB11 wB12 wB2 wB3 wB4 wB51 wB52 wB53 wB6 wB71 wB72 wB81 wB82 wB83 wB9]
        resonancewB = [0.009, 0.010, 0.0005, 0.0005, 0.001, 0.0005, 0.021, 0.015, 0.043, 0.0005, 0.130, 0.024, 0.0005, 0.036, 3.1] * u.G #resonance width
        
    else: # old values
        
        if ABKG_choice == 1:
            a_bkg = 87.2 * ac.a0
        elif ABKG_choice == 2:
            a_bkg = 95.2 * ac.a0
        elif ABKG_choice == 3:
            a_bkg = 79.2 * ac.a0

        #FR resonances
        #[B1 B2 B3 B4 B5 B6 B7 B8]
        resonanceB  = [1.298, 2.802, 3.370, 5.092, 7.154, 2.592, 2.338, 2.177] * u.G #resonance position
        #[wB1 wB2 wB3 wB4 wB5 wB6 wB7 wB8]
        resonancewB = [0.018, 0.047, 0.048, 0.145, 0.020, 0.008, 0.001, 0.001] * u.G #resonance width
        
    #Get scattering length
    np.seterr(divide='ignore')
    a_s = a_bkg * np.prod([(1 - resonancewB[j] / (B - resonanceB[j])) for j in range(len(resonanceB))])
    return a_s, a_bkg

def dipolarLength(mu = DY_DIPOLE_MOMENT, m = DY_MASS):
    return (m * ac.mu0 * mu**2) / (12 * np.pi * ac.hbar**2)

def scatteringCrossSection(B):
    return 8 * np.pi * scatteringLength(B)[0]**2 + ((32*np.pi)/45) * dipolarLength()**2

def calculateElasticCollisionRate(w_x, w_z, Power, N, T, B): #For a 3D Harmonic Trap
    return (particleDensity(w_x, w_z, Power, N, T) * scatteringCrossSection(B) * meanThermalVelocity(T) / (2 * np.sqrt(2))).decompose()

def calculatePSD(w_x, w_z, Power, N, T):
    return (particleDensity(w_x, w_z, Power, N, T) * thermaldeBroglieWavelength(T)**3).decompose()

def convert_modulation_depth_to_alpha(modulation_depth):
    fin_mod_dep = [0, 0.5, 0.3, 0.7, 0.9, 0.8, 1.0, 0.6, 0.4, 0.2, 0.1]
    fx = [3.135, 0.28, 0.690, 0.152, 0.102, 0.127, 0.099, 0.205, 0.404, 1.441, 2.813]
    dfx = [0.016, 0.006, 0.005, 0.006, 0.003, 0.002, 0.002,0.002, 0.003, 0.006, 0.024]
    fz = [2.746, 1.278, 1.719, 1.058, 0.923, 0.994, 0.911, 1.157, 1.446, 2.191, 2.643]
    dfz = [0.014, 0.007, 0.009, 0.007, 0.005, 0.004, 0.004, 0.005, 0.007, 0.009, 0.033]

    alpha_x =  [(fx[0]/x)**(2/3) for x in fx]
    dalpha_x = [alpha_x[i] * np.sqrt((dfx[0]/fx[0])**2 + (dfx[i]/fx[i])**2) for i in range(len(fx))]
    alpha_z = [(fz[0]/z)**2 for z in fz]
    dalpha_z = [alpha_z[i] * np.sqrt((dfz[0]/fz[0])**2 + (dfz[i]/fz[i])**2) for i in range(len(fz))]

    avg_alpha = [(g + h) / 2 for g, h in zip(alpha_x, alpha_z)]
    sorted_fin_mod_dep, sorted_avg_alpha = zip(*sorted(zip(fin_mod_dep, avg_alpha)))

    f = interpolate.interp1d(sorted_fin_mod_dep, sorted_avg_alpha)

    return f(modulation_depth), fin_mod_dep, alpha_x, alpha_z, dalpha_x, dalpha_z

def convert_modulation_depth_to_temperature(modulation_depth):
    fin_mod_dep = [1.0, 0.8, 0.6, 0.4, 0.2, 0.0, 0.9, 0.7, 0.5, 0.3, 0.1]
    T_x = [22.1, 27.9, 31.7, 42.2, 98.8, 145.8, 27.9, 33.8, 42.4, 61.9, 136.1]
    dT_x = [1.7, 2.6, 2.4, 3.7, 1.1, 0.6, 2.2, 3.2, 1.7, 2.2, 1.2]
    T_y = [13.13, 14.75, 18.44, 26.31, 52.55, 92.9, 13.54, 16.11, 21.15, 35.81, 85.8]
    dT_y = [0.05, 0.05, 0.07, 0.16, 0.28, 0.7, 0.04, 0.07, 0.10, 0.21, 0.8]

    avg_T = [(g + h) / 2 for g, h in zip(T_x, T_y)]
    sorted_fin_mod_dep, sorted_avg_T = zip(*sorted(zip(fin_mod_dep, avg_T)))

    f = interpolate.interp1d(sorted_fin_mod_dep, sorted_avg_T)

    return f(modulation_depth), fin_mod_dep, T_x, T_y, dT_x, dT_y

#####################################################################
# 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, beam_1_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, beam_1_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

#####################################################################
# COMPUTE/EXTRACT TRAP POTENTIAL AND PARAMETERS                     #
#####################################################################

def trap_depth(w_1, w_2, P, alpha = DY_POLARIZABILITY):
    return 2*P/(np.pi*w_1*w_2) * (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)

def calculateTrapFrequency(w_x, w_z, Power, dir = 'x', m = DY_MASS):
    TrapDepth = trap_depth(w_x, w_z, Power)
    TrapFrequency = np.nan
    if dir == 'x':
        TrapFrequency = ((1/(2 * np.pi)) * np.sqrt(4 * TrapDepth / (m*w_x**2))).decompose()
    elif dir == 'y':
        zReff = np.sqrt(2) * z_R(w_x, 1.064*u.um) * z_R(w_z, 1.064*u.um) / np.sqrt(z_R(w_x, 1.064*u.um)**2 + z_R(w_z, 1.064*u.um)**2)
        TrapFrequency = ((1/(2 * np.pi)) * np.sqrt(2 * TrapDepth/ (m*zReff**2))).decompose()
    elif dir == 'z':
        TrapFrequency = ((1/(2 * np.pi)) * np.sqrt(4 * TrapDepth/ (m*w_z**2))).decompose()
    return round(TrapFrequency.value, 2)*u.Hz

def calculateCrossedBeamTrapFrequency(delta, Waists, Powers, dir = 'x', m = DY_MASS, wavelength=1.064*u.um):
    
    TrapDepth_1 = trap_depth(Waists[0][0], Waists[1][0], Powers[0])
    TrapDepth_2 = trap_depth(Waists[0][1], Waists[1][1], Powers[1])
    
    w_x1 = Waists[0][0]
    w_z1 = Waists[1][0]

    w_y2 = Waists[0][1]
    w_z2 = Waists[1][1] 
        
    zReff_1 = np.sqrt(2) * z_R(w_x1, wavelength) * z_R(w_z1, wavelength) / np.sqrt(z_R(w_x1, wavelength)**2 + z_R(w_z1, wavelength)**2)
    zReff_2 = np.sqrt(2) * z_R(w_y2, wavelength) * z_R(w_z2, wavelength) / np.sqrt(z_R(w_y2, wavelength)**2 + z_R(w_z2, wavelength)**2)
    
    wy2alpha = np.sqrt(((np.cos(np.radians(90 - delta)) / w_y2)**2 + (np.sin(np.radians(90 - delta)) / (2 * zReff_2))**2)**(-1))
    zR2alpha = np.sqrt(((np.sin(np.radians(90 - delta)) / w_y2)**2 + (np.cos(np.radians(90 - delta)) / (2 * zReff_2))**2)**(-1))    
   
    TrapFrequency = np.nan
    if dir == 'x':
        v_1 = (1/(2 * np.pi)) * np.sqrt(4/m * (TrapDepth_1 / w_x1**2))
        v_2 = (1/(2 * np.pi)) * np.sqrt(4/m * (TrapDepth_2 / zR2alpha**2))
        TrapFrequency = (np.sqrt(v_1**2 + v_2**2)).decompose()
    elif dir == 'y':
        v_1 = (1/(2 * np.pi)) * np.sqrt(1/m * (2 * TrapDepth_1/ zReff_1**2))
        v_2 = (1/(2 * np.pi)) * np.sqrt(1/m * (4 * TrapDepth_2/ wy2alpha**2))
        TrapFrequency = (np.sqrt(v_1**2 + v_2**2)).decompose()
    elif dir == 'z':
        v_1 = (1/(2 * np.pi)) * np.sqrt(4/m * (TrapDepth_1 / w_z1**2))
        v_2 = (1/(2 * np.pi)) * np.sqrt(4/m * (TrapDepth_2 / w_z2**2))
        TrapFrequency = (np.sqrt(v_1**2 + v_2**2)).decompose()

    return round(TrapFrequency.value, 2)*u.Hz

def extractTrapFrequency(Positions, TrappingPotential, axis):
    tmp_pos = Positions[axis, :]
    tmp_pot = TrappingPotential[axis]
    center_idx = np.argmin(tmp_pot)
    lb = int(round(center_idx - len(tmp_pot)/200, 1))
    ub = int(round(center_idx + len(tmp_pot)/200, 1))
    xdata = tmp_pos[lb:ub]
    Potential = tmp_pot[lb:ub]
    p0 = [1e3, tmp_pos[center_idx].value, -100]
    popt, pcov = curve_fit(harmonic_potential, xdata, Potential, p0)
    v = popt[0]
    dv = pcov[0][0]**0.5 
    return v, dv, popt, pcov

def computeTrapPotential(w_x, w_z, Power, options):

    axis = options['axis']
    extent = options['extent']
    gravity = options['gravity']
    astigmatism = options['astigmatism']
    modulation = options['modulation']
    crossed = options['crossed']

    if modulation:
        aspect_ratio = options['aspect_ratio']
        current_ar = w_x/w_z
        w_x = w_x * (aspect_ratio / current_ar)

    TrappingPotential = []
    TrapDepth = trap_depth(w_x, w_z, Power) 
    IdealTrapDepthInKelvin = (TrapDepth/ac.k_B).to(u.uK)
    
    projection_axis = np.array([0, 1, 0]) # default

    if axis == 0:
        projection_axis = np.array([1, 0, 0]) # radial direction (X-axis)
    
    elif axis == 1:
        projection_axis = np.array([0, 1, 0]) # propagation direction (Y-axis)
    
    elif axis == 2:
        projection_axis = np.array([0, 0, 1]) # vertical direction (Z-axis)
    
    else:
        projection_axis = np.array([1, 1, 1]) # vertical direction (Z-axis)
    
    x_Positions      = np.arange(-extent, extent, 0.05)*u.um 
    y_Positions      = np.arange(-extent, extent, 0.05)*u.um
    z_Positions      = np.arange(-extent, extent, 0.05)*u.um 
    Positions        = np.vstack((x_Positions, y_Positions, z_Positions)) * projection_axis[:, np.newaxis]

    if not crossed:
        IdealTrappingPotential        = single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power) 
        IdealTrappingPotential        = IdealTrappingPotential * (np.ones((3, len(IdealTrappingPotential))) * projection_axis[:, np.newaxis])
        IdealTrappingPotential        = (IdealTrappingPotential/ac.k_B).to(u.uK) 

        if gravity and not astigmatism:
            # Influence of Gravity
            m = DY_MASS
            gravity_axis             = np.array([0, 0, 1]) 
            tilt_gravity             = options['tilt_gravity']
            theta                    = options['theta']
            tilt_axis                = options['tilt_axis']
            if tilt_gravity:
                R = rotation_matrix(tilt_axis, np.radians(theta))
                gravity_axis = np.dot(R, gravity_axis)
            gravity_axis_positions   = np.vstack((x_Positions, y_Positions, z_Positions)) * gravity_axis[:, np.newaxis]
            TrappingPotential        = single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power)
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis]) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
            
        elif not gravity and astigmatism:
            # Influence of Astigmatism
            foci_disp_single = options['foci_disp_single']
            TrappingPotential        = astigmatic_single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power, del_y = foci_disp_single)
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis])
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        elif gravity and astigmatism:
            # Influence of Gravity and Astigmatism
            m = DY_MASS
            gravity_axis             = np.array([0, 0, 1]) 
            tilt_gravity             = options['tilt_gravity']
            theta                    = options['theta']
            tilt_axis                = options['tilt_axis']
            foci_disp_single         = options['foci_disp_single']
            if tilt_gravity:
                R = rotation_matrix(tilt_axis, np.radians(theta))
                gravity_axis = np.dot(R, gravity_axis)
            gravity_axis_positions   = np.vstack((x_Positions, y_Positions, z_Positions)) * gravity_axis[:, np.newaxis]
            TrappingPotential        = astigmatic_single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power, del_y = foci_disp_single)
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis]) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        else:
            TrappingPotential = IdealTrappingPotential
    
        infls = np.where(np.diff(np.sign(np.gradient(np.gradient(TrappingPotential[axis].value)))))[0]
        
        try:
            if TrappingPotential[axis][0] > TrappingPotential[axis][-1]:
                EffectiveTrapDepthInKelvin = max(TrappingPotential[axis][infls[1]:-1]) - min(TrappingPotential[axis][infls[0]:infls[1]])
            elif TrappingPotential[axis][0] < TrappingPotential[axis][-1]:
                EffectiveTrapDepthInKelvin = max(TrappingPotential[axis][0:infls[0]]) - min(TrappingPotential[axis][infls[0]:infls[1]])            
            else:
                EffectiveTrapDepthInKelvin = IdealTrapDepthInKelvin
        except:
            EffectiveTrapDepthInKelvin = np.nan

        TrapDepthsInKelvin = [IdealTrapDepthInKelvin, EffectiveTrapDepthInKelvin]

        v_x = calculateTrapFrequency(w_x, w_z, Power, dir = 'x')
        v_y = calculateTrapFrequency(w_x, w_z, Power, dir = 'y')
        v_z = calculateTrapFrequency(w_x, w_z, Power, dir = 'z')
        CalculatedTrapFrequencies = [v_x, v_y, v_z]

        v, dv, popt, pcov = extractTrapFrequency(Positions, IdealTrappingPotential, axis)
        if np.isinf(v):
            v = np.nan
        if np.isinf(dv):
            dv = np.nan

        IdealTrapFrequency = [v, dv]
        
        if options['extract_trap_frequencies']:
            v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, axis)
            if np.isinf(v):
                v = np.nan
            if np.isinf(dv):
                dv = np.nan
            TrapFrequency = [v, dv]
            ExtractedTrapFrequencies = [IdealTrapFrequency, TrapFrequency]
        else:
            ExtractedTrapFrequencies = [IdealTrapFrequency]
        
        return Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies

    else:
        waists = np.vstack((np.asarray([w_x[0].value, w_z[0].value])*u.um, np.asarray([w_x[1].value, w_z[1].value])*u.um))
        IdealTrappingPotential        = crossed_beam_potential(Positions, waists, Power, options) 
        IdealTrappingPotential        = IdealTrappingPotential * (np.ones((3, len(IdealTrappingPotential))) * projection_axis[:, np.newaxis])
        IdealTrappingPotential        = (IdealTrappingPotential/ac.k_B).to(u.uK) 

        if gravity and not astigmatism:
            # Influence of Gravity
            m = DY_MASS
            gravity_axis             = np.array([0, 0, 1]) 
            tilt_gravity             = options['tilt_gravity']
            theta                    = options['theta']
            tilt_axis                = options['tilt_axis']
            if tilt_gravity:
                R = rotation_matrix(tilt_axis, np.radians(theta))
                gravity_axis = np.dot(R, gravity_axis)
            gravity_axis_positions   = np.vstack((x_Positions, y_Positions, z_Positions)) * gravity_axis[:, np.newaxis]
            TrappingPotential        = crossed_beam_potential(Positions, waists, Power, options) 
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis]) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
            
        elif not gravity and astigmatism:
            # Influence of Astigmatism
            TrappingPotential        = astigmatic_crossed_beam_potential(Positions, waists, Power, options)
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis])
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        elif gravity and astigmatism:
            # Influence of Gravity and Astigmatism
            m = DY_MASS
            gravity_axis             = np.array([0, 0, 1]) 
            tilt_gravity             = options['tilt_gravity']
            theta                    = options['theta']
            tilt_axis                = options['tilt_axis']
            if tilt_gravity:
                R = rotation_matrix(tilt_axis, np.radians(theta))
                gravity_axis = np.dot(R, gravity_axis)
            gravity_axis_positions   = np.vstack((x_Positions, y_Positions, z_Positions)) * gravity_axis[:, np.newaxis]
            TrappingPotential        = astigmatic_crossed_beam_potential(Positions, waists, Power, options)
            TrappingPotential        = TrappingPotential * (np.ones((3, len(TrappingPotential))) * projection_axis[:, np.newaxis]) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        else:
            TrappingPotential = IdealTrappingPotential

        infls = np.where(np.diff(np.sign(np.gradient(np.gradient(TrappingPotential[axis].value)))))[0]
        
        try:
            if TrappingPotential[axis][0] > TrappingPotential[axis][-1]:
                EffectiveTrapDepthInKelvin = max(TrappingPotential[axis][infls[1]:-1]) - min(TrappingPotential[axis][infls[0]:infls[1]])
            elif TrappingPotential[axis][0] < TrappingPotential[axis][-1]:
                EffectiveTrapDepthInKelvin = max(TrappingPotential[axis][0:infls[0]]) - min(TrappingPotential[axis][infls[0]:infls[1]])            
            else:
                EffectiveTrapDepthInKelvin = IdealTrapDepthInKelvin
        except:
            EffectiveTrapDepthInKelvin = np.nan

        TrapDepthsInKelvin = [IdealTrapDepthInKelvin, EffectiveTrapDepthInKelvin]

        v_x = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Power, dir = 'x')
        v_y = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Power, dir = 'y')
        v_z = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Power, dir = 'z')
        CalculatedTrapFrequencies = [v_x, v_y, v_z]

        v, dv, popt, pcov = extractTrapFrequency(Positions, IdealTrappingPotential, axis)
        if np.isinf(v):
            v = np.nan
        if np.isinf(dv):
            dv = np.nan

        IdealTrapFrequency = [v, dv]
        
        if options['extract_trap_frequencies']:
            v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, axis)
            if np.isinf(v):
                v = np.nan
            if np.isinf(dv):
                dv = np.nan
            TrapFrequency = [v, dv]
            ExtractedTrapFrequencies = [IdealTrapFrequency, TrapFrequency]
        else:
            ExtractedTrapFrequencies = [IdealTrapFrequency]

        return Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies

def extractWaist(Positions, TrappingPotential):
    tmp_pos     = Positions.value
    tmp_pot     = TrappingPotential.value
    center_idx  = np.argmin(tmp_pot)
    
    TrapMinimum = tmp_pot[center_idx]
    TrapCenter  = tmp_pos[center_idx]

    lb = int(round(center_idx - len(tmp_pot)/30, 1))
    ub = int(round(center_idx + len(tmp_pot)/30, 1))
    xdata = tmp_pos[lb:ub]
    Potential = tmp_pot[lb:ub]

    p0=[TrapMinimum, 30, TrapCenter, 0]
    popt, pcov = curve_fit(gaussian_potential, xdata, Potential, p0)
    return popt, pcov

def computeIntensityProfileAndPotentials(Power, waists, wavelength, options, alpha = DY_POLARIZABILITY):
    w_x = waists[0]
    w_z = waists[1]
    extent = options['extent']
    modulation = options['modulation']
    mod_func = options['modulation_function']

    if not modulation:
        extent = 50
        x_Positions      = np.arange(-extent, extent, 1)*u.um 
        y_Positions      = np.arange(-extent, extent, 1)*u.um
        z_Positions      = np.arange(-extent, extent, 1)*u.um 
        
        idx = np.where(y_Positions==0)[0][0]

        wavelength = 1.064*u.um
        
        xm,ym,zm = np.meshgrid(x_Positions, y_Positions, z_Positions, sparse=True, indexing='ij')

        ## Single Gaussian Beam
        A = 2*Power/(np.pi*w(ym, w_x, wavelength)*w(ym, w_z, wavelength))
        intensity_profile = A * np.exp(-2 * ((xm/w(ym, w_x, wavelength))**2 + (zm/w(ym, w_z, wavelength))**2))
        I = intensity_profile[:, idx, :].to(u.MW/(u.cm*u.cm))
        
        U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
        U = - U_tilde * I
        U = (U/ac.k_B).to(u.uK)
        
        return [x_Positions, z_Positions], [w_x.value, 0, w_z.value, 0], I, U, [0, 0, 0, 0]

    else:
        mod_amp = options['modulation_amplitude']
        x_Positions      = np.arange(-extent, extent, 1)*u.um 
        y_Positions      = np.arange(-extent, extent, 1)*u.um
        z_Positions      = np.arange(-extent, extent, 1)*u.um 

        mod_amp   = mod_amp * w_x
        n_points  = len(x_Positions)
        dx, xmod_Positions = modulation_function(mod_amp, n_points, func = mod_func)
        
        idx = np.where(y_Positions==0)[0][0]

        xm,ym,zm,xmodm = np.meshgrid(x_Positions, y_Positions, z_Positions, xmod_Positions, sparse=True, indexing='ij')

        ## Single Modulated Gaussian Beam
        A = 2*Power/(np.pi*w(y_Positions[idx] , w_x, wavelength)*w(y_Positions[idx], w_z, wavelength))
        intensity_profile = A * 1/(2*mod_amp) * np.trapz(np.exp(-2 * (((xmodm - xm)/w(ym, w_x, wavelength))**2 + (zm/w(ym, w_z, wavelength))**2)), dx = dx, axis = -1)
        I = intensity_profile[:, idx, :].to(u.MW/(u.cm*u.cm))
        
        U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
        U = - U_tilde * I
        U = (U/ac.k_B).to(u.uK)

        poptx, pcovx = extractWaist(x_Positions, U[:, np.where(z_Positions==0)[0][0]])
        poptz, pcovz = extractWaist(z_Positions, U[np.where(x_Positions==0)[0][0], :])
        
        extracted_waist_x = poptx[1]
        dextracted_waist_x = pcovx[1][1]**0.5
        extracted_waist_z = poptz[1]
        dextracted_waist_z = pcovz[1][1]**0.5
        
        return [x_Positions, z_Positions], [extracted_waist_x, dextracted_waist_x, extracted_waist_z, dextracted_waist_z], I, U, [poptx, pcovx, poptz, pcovz]

#####################################################################
# PLOTTING                                                          #
#####################################################################

def generate_label(v, dv):
    unit = 'Hz'
    if v <= 0.0:
        v        = np.nan
        dv       = np.nan
        unit = 'Hz'
    elif v > 0.0 and orderOfMagnitude(v) > 2:
        v    = v / 1e3 # in kHz
        dv   = dv / 1e3 # in kHz
        unit = 'kHz'
    tf_label = '\u03BD = %.1f \u00B1 %.2f %s'% tuple([v,dv,unit])
    return tf_label

def plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, axis, popt, pcov):
    v = popt[0]
    dv = pcov[0][0]**0.5
    happrox = harmonic_potential(Positions[axis, :].value, *popt)
    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(121)
    ax.set_title('Fit to Potential')
    plt.plot(Positions[axis, :].value, happrox, '-r', label = '\u03BD = %.1f \u00B1 %.2f Hz'% tuple([v,dv]))
    plt.plot(Positions[axis, :], TrappingPotential[axis], 'ob', label = 'Gaussian Potential')
    plt.xlabel('Distance (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')
    plt.ylim([-TrapDepthsInKelvin[0].value, max(TrappingPotential[axis].value)])
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})

    bx = fig.add_subplot(122)
    bx.set_title('Fit Residuals')
    plt.plot(Positions[axis, :].value, TrappingPotential[axis].value - happrox, 'ob')
    plt.xlabel('Distance (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('$U_{trap} - U_{Harmonic}$', fontsize= 12, fontweight='bold')
    plt.xlim([-10, 10])
    plt.ylim([-1e-2, 1e-2])
    plt.grid(visible=1)
    plt.tight_layout()
    plt.show()

def plotGaussianFit(Positions, TrappingPotential, popt, pcov):
    extracted_waist = popt[1]
    dextracted_waist = pcov[1][1]**0.5
    gapprox = gaussian_potential(Positions, *popt)
    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(121)
    ax.set_title('Fit to Potential')
    plt.plot(Positions, gapprox, '-r', label = 'waist = %.1f \u00B1 %.2f um'% tuple([extracted_waist,dextracted_waist]))
    plt.plot(Positions, TrappingPotential, 'ob', label = 'Gaussian Potential')
    plt.xlabel('Distance (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')
    plt.ylim([min(TrappingPotential), max(TrappingPotential)])
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})

    bx = fig.add_subplot(122)
    bx.set_title('Fit Residuals')
    plt.plot(Positions, TrappingPotential - gapprox, 'ob')
    plt.xlabel('Distance (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('$U_{trap} - U_{Gaussian}$', fontsize= 12, fontweight='bold')
    plt.xlim([-10, 10])
    plt.ylim([-1, 1])
    plt.grid(visible=1)
    plt.tight_layout()
    plt.show()

def plotPotential(Positions, ComputedPotentials, options, Params = [], listToIterateOver = [], save = False):

    axis = options['axis']

    plt.figure(figsize=(9, 7))
    for i in range(np.size(ComputedPotentials, 0)):
        
        if i % 2 == 0:
            j = int(i / 2)
        else:
            j = int((i - 1) / 2)
        
        IdealTrapDepthInKelvin = Params[j][0][0]
        EffectiveTrapDepthInKelvin = Params[j][0][1]
        
        idealv = Params[j][2][0][0]
        idealdv = Params[j][2][0][1]
        
        if options['extract_trap_frequencies']:
            v = Params[j][2][1][0]
            dv = Params[j][2][1][1]
        else:
            v = np.nan
            dv = np.nan
        
        if listToIterateOver:
            if np.size(ComputedPotentials, 0) == len(listToIterateOver):
                plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'Trap Depth = ' + str(round(EffectiveTrapDepthInKelvin.value, 2)) + ' ' + str(EffectiveTrapDepthInKelvin.unit) + '; ' + generate_label(v, dv)) 
            else:
                if i % 2 == 0:
                    plt.plot(Positions[axis], ComputedPotentials[i][axis], '--', label = 'Trap Depth = ' + str(round(IdealTrapDepthInKelvin.value, 2)) + ' ' + str(IdealTrapDepthInKelvin.unit) + '; ' + generate_label(idealv, idealdv)) 
                elif i % 2 != 0:
                    plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'Effective Trap Depth = ' + str(round(EffectiveTrapDepthInKelvin.value, 2)) + ' ' + str(EffectiveTrapDepthInKelvin.unit) + '; ' + generate_label(v, dv))
        else:
            if i % 2 == 0:
                plt.plot(Positions[axis], ComputedPotentials[i][axis], '--', label = 'Trap Depth = ' + str(round(IdealTrapDepthInKelvin.value, 2)) + ' ' + str(IdealTrapDepthInKelvin.unit) + '; ' + generate_label(idealv, idealdv))
            elif i % 2 != 0:
                plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'Effective Trap Depth = ' + str(round(EffectiveTrapDepthInKelvin.value, 2)) + ' ' + str(EffectiveTrapDepthInKelvin.unit) + '; ' + generate_label(v, dv))
    if axis == 0:
        dir = 'X - Horizontal'
    elif axis == 1:
        dir = 'Y - Propagation'
    else:
        dir = 'Z - Vertical'
    
    plt.ylim(top = 0)
    plt.xlabel(dir + ' Direction (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(loc=3, prop={'size': 12, 'weight': 'bold'})
    if save:
        plt.savefig('pot_' + dir + '.png')
    plt.show()

def plotIntensityProfileAndPotentials(positions, waists, I, U):
    
    x_Positions = positions[0]
    z_Positions = positions[1]
    
    w_x = waists[0]
    dw_x = waists[1]
    w_z = waists[2]
    dw_x = waists[3]

    ar = w_x/w_z
    dar = ar * np.sqrt((dw_x/w_x)**2 + (dw_x/w_z)**2)
    
    fig = plt.figure(figsize=(12, 6))
    ax = fig.add_subplot(121)
    ax.set_title('Intensity Profile ($MW/cm^2$)\n Aspect Ratio = %.2f \u00B1 %.2f um'% tuple([ar,dar]))
    im = plt.imshow(np.transpose(I.value), cmap="coolwarm", extent=[np.min(x_Positions.value), np.max(x_Positions.value), np.min(z_Positions.value), np.max(z_Positions.value)])
    plt.xlabel('X - Horizontal (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('Z - Vertical (um)', fontsize= 12, fontweight='bold')
    ax.set_aspect('equal')
    fig.colorbar(im, fraction=0.046, pad=0.04, orientation='vertical')
    
    bx = fig.add_subplot(122)
    bx.set_title('Trap Potential')
    plt.plot(x_Positions, U[:, np.where(z_Positions==0)[0][0]], label = 'X - Horizontal')
    plt.plot(z_Positions, U[np.where(x_Positions==0)[0][0], :], label = 'Z - Vertical')
    plt.ylim(top = 0)
    plt.xlabel('Extent (um)', fontsize= 12, fontweight='bold')
    plt.ylabel('Depth (uK)', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show()

def plotAlphas():

    modulation_depth = np.arange(0, 1.1, 0.1)
    Alphas, fin_mod_dep, alpha_x, alpha_y, dalpha_x, dalpha_y = convert_modulation_depth_to_alpha(modulation_depth)   

    plt.figure()
    plt.errorbar(fin_mod_dep, alpha_x, yerr = dalpha_x, fmt= 'ob', label = 'From Horz TF', markersize=5, capsize=5)
    plt.errorbar(fin_mod_dep, alpha_y, yerr = dalpha_y, fmt= 'or', label = 'From Vert TF', markersize=5, capsize=5)
    plt.plot(modulation_depth, Alphas, '--g')
    plt.xlabel('Modulation depth', fontsize= 12, fontweight='bold')
    plt.ylabel('$\\alpha$', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show()

def plotTemperatures(w_x, w_z, plot_against_mod_depth = True):

    modulation_depth = np.arange(0, 1.1, 0.1)
    w_xs = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]
    new_aspect_ratio  = w_xs / w_z
    Temperatures, fin_mod_dep, T_x, T_y, dT_x, dT_y = convert_modulation_depth_to_temperature(modulation_depth)  
    measured_aspect_ratio = (w_x * convert_modulation_depth_to_alpha(fin_mod_dep)[0]) / w_z

    plt.figure()
    if plot_against_mod_depth:
        plt.errorbar(fin_mod_dep, T_x, yerr = dT_x, fmt= 'ob', label = 'Horz direction', markersize=5, capsize=5)
        plt.errorbar(fin_mod_dep, T_y, yerr = dT_y, fmt= 'or', label = 'Vert direction', markersize=5, capsize=5)
        plt.plot(modulation_depth, Temperatures, '--g')
        xlabel = 'Modulation depth'
    else:
        plt.errorbar(measured_aspect_ratio, T_x, yerr = dT_x, fmt= 'ob', label = 'Horz direction', markersize=5, capsize=5)
        plt.errorbar(measured_aspect_ratio, T_y, yerr = dT_y, fmt= 'or', label = 'Vert direction', markersize=5, capsize=5)
        plt.plot(new_aspect_ratio, Temperatures, '--g')
        xlabel = 'Aspect Ratio'

    plt.xlabel(xlabel, fontsize= 12, fontweight='bold')
    plt.ylabel('Temperature (uK)', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show()

def plotTrapFrequencies(v_x, v_y, v_z, modulation_depth, new_aspect_ratio, plot_against_mod_depth = True):
    fig, ax3 = plt.subplots(figsize=(8, 6))
    
    if plot_against_mod_depth:
        ln1 = ax3.plot(modulation_depth, v_x, '-or', label = 'v_x')
        ln2 = ax3.plot(modulation_depth, v_z, '-^b', label = 'v_z')
        ax4 = ax3.twinx()
        ln3 = ax4.plot(modulation_depth, v_y, '-*g', label = 'v_y')
        xlabel = 'Modulation depth'
    else:
        ln1 = ax3.plot(new_aspect_ratio, v_x, '-or', label = 'v_x')
        ln2 = ax3.plot(new_aspect_ratio, v_z, '-^b', label = 'v_z')
        ax4 = ax3.twinx()
        ln3 = ax4.plot(new_aspect_ratio, v_y, '-*g', label = 'v_y')
        xlabel = 'Aspect Ratio'

    ax3.set_xlabel(xlabel, fontsize= 12, fontweight='bold')
    ax3.set_ylabel('Trap Frequency (Hz)', fontsize= 12, fontweight='bold')
    ax3.tick_params(axis="y", labelcolor='b')
    ax4.set_ylabel('Trap Frequency (Hz)', fontsize= 12, fontweight='bold')
    ax4.tick_params(axis="y", labelcolor='g')
    plt.tight_layout()
    plt.grid(visible=1)
    lns = ln1+ln2+ln3
    labs = [l.get_label() for l in lns]
    ax3.legend(lns, labs, prop={'size': 12, 'weight': 'bold'})
    plt.show()

def plotMeasuredTrapFrequencies(fx, dfx, fy, dfy, fz, dfz,  modulation_depth_radial, modulation_depth_axial, w_x, w_z, plot_against_mod_depth = True):
    
    alpha_x =  [(fx[0]/x)**(2/3) for x in fx]
    dalpha_x = [alpha_x[i] * np.sqrt((dfx[0]/fx[0])**2 + (dfx[i]/fx[i])**2) for i in range(len(fx))]
    alpha_y = [(fy[0]/y)**2 for y in fy]
    dalpha_y = [alpha_y[i] * np.sqrt((dfy[0]/fy[0])**2 + (dfy[i]/fy[i])**2) for i in range(len(fy))]

    avg_alpha = [(g + h) / 2 for g, h in zip(alpha_x, alpha_y)]
    new_aspect_ratio = (w_x * avg_alpha) / w_z
    
    
    if plot_against_mod_depth:
        fig, ax1 = plt.subplots(figsize=(8, 6))
        ax2 = ax1.twinx()
        ax1.errorbar(modulation_depth_radial, fx, yerr = dfx, fmt= 'or', label = 'v_x', markersize=5, capsize=5)
        ax2.errorbar(modulation_depth_axial, fy, yerr = dfy, fmt= '*g', label = 'v_y', markersize=5, capsize=5)
        ax1.errorbar(modulation_depth_radial, fz, yerr = dfz, fmt= '^b', label = 'v_z', markersize=5, capsize=5)
        ax1.set_xlabel('Modulation depth', fontsize= 12, fontweight='bold')
        ax1.set_ylabel('Trap Frequency (kHz)', fontsize= 12, fontweight='bold')
        ax1.tick_params(axis="y", labelcolor='b')
        ax2.set_ylabel('Trap Frequency (Hz)', fontsize= 12, fontweight='bold')
        ax2.tick_params(axis="y", labelcolor='g')
        h1, l1 = ax1.get_legend_handles_labels()
        h2, l2 = ax2.get_legend_handles_labels()
        ax1.legend(h1+h2, l1+l2, loc=0, prop={'size': 12, 'weight': 'bold'})
    else:
        plt.figure()
        plt.errorbar(new_aspect_ratio, fx, yerr = dfx, fmt= 'or', label = 'v_x', markersize=5, capsize=5)
        plt.errorbar(new_aspect_ratio, fz, yerr = dfz, fmt= '^b', label = 'v_z', markersize=5, capsize=5)
        plt.xlabel('Aspect Ratio', fontsize= 12, fontweight='bold')
        plt.ylabel('Trap Frequency (kHz)', fontsize= 12, fontweight='bold')
        plt.legend(prop={'size': 12, 'weight': 'bold'})

    plt.tight_layout()
    plt.grid(visible=1)
    plt.show() 

def plotRatioOfTrapFrequencies(fx, fy, fz, dfx, dfy, dfz, v_x, v_y, v_z, modulation_depth, w_x, w_z, plot_against_mod_depth = True):
    
    w_xs = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]
    new_aspect_ratio  = w_xs / w_z

    plt.figure()
    
    if plot_against_mod_depth:
        plt.errorbar(modulation_depth, fx/v_x, yerr = dfx/v_x, fmt= 'or', label = 'b/w horz TF', markersize=5, capsize=5)
        plt.errorbar(modulation_depth, fy/v_y, yerr = dfy/v_y, fmt= '*g', label = 'b/w axial TF', markersize=5, capsize=5)
        plt.errorbar(modulation_depth, fz/v_z, yerr = dfz/v_z, fmt= '^b', label = 'b/w vert TF', markersize=5, capsize=5)
        xlabel = 'Modulation depth'
    else:
        plt.errorbar(new_aspect_ratio, fx/v_x, yerr = dfx/v_x, fmt= 'or', label = 'b/w horz TF', markersize=5, capsize=5)
        plt.errorbar(new_aspect_ratio, fy/v_y, yerr = dfy/v_y, fmt= '*g', label = 'b/w axial TF', markersize=5, capsize=5)
        plt.errorbar(new_aspect_ratio, fz/v_z, yerr = dfz/v_z, fmt= '^b', label = 'b/w vert TF', markersize=5, capsize=5)
        xlabel = 'Aspect Ratio'

    plt.xlabel(xlabel, fontsize= 12, fontweight='bold')
    plt.ylabel('Ratio', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show() 

def plotScatteringLengths(B_range = [0, 2.59]):
    BField = np.arange(B_range[0], B_range[1], 1e-3) * u.G
    a_s_array = np.zeros(len(BField)) * ac.a0
    for idx in range(len(BField)):
        a_s_array[idx], a_bkg = scatteringLength(BField[idx])
    rmelmIdx = [i for i, x in enumerate(np.isinf(a_s_array.value)) if x] 
    for x in rmelmIdx:
        a_s_array[x-1] = np.inf * ac.a0
    
    plt.figure(figsize=(9, 7))
    plt.plot(BField, a_s_array/ac.a0, '-b')
    plt.axhline(y = a_bkg/ac.a0, color = 'r', linestyle = '--')
    plt.text(min(BField.value) + 0.5, (a_bkg/ac.a0).value + 1, '$a_{bkg}$ = %.2f a0' %((a_bkg/ac.a0).value), fontsize=14, fontweight='bold')
    plt.xlim([min(BField.value), max(BField.value)])
    plt.ylim([65, 125])
    plt.xlabel('B field (G)', fontsize= 12, fontweight='bold')
    plt.ylabel('Scattering length (a0)', fontsize= 12, fontweight='bold')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.show()

def plotCollisionRatesAndPSD(Gamma_elastic, PSD, modulation_depth, new_aspect_ratio, plot_against_mod_depth = True):
    fig, ax1 = plt.subplots(figsize=(8, 6))
    ax2 = ax1.twinx()

    if plot_against_mod_depth:
        ax1.plot(modulation_depth, Gamma_elastic, '-ob')
        ax2.plot(modulation_depth, PSD, '-*r')
        ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
        xlabel = 'Modulation depth'
    else:
        ax1.plot(new_aspect_ratio, Gamma_elastic, '-ob')
        ax2.plot(new_aspect_ratio, PSD, '-*r')
        ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
        xlabel = 'Aspect Ratio'

    ax1.set_xlabel(xlabel, fontsize= 12, fontweight='bold')
    ax1.set_ylabel('Elastic Collision Rate', fontsize= 12, fontweight='bold')
    ax1.tick_params(axis="y", labelcolor='b')
    ax2.set_ylabel('Phase Space Density', fontsize= 12, fontweight='bold')
    ax2.tick_params(axis="y", labelcolor='r')
    plt.tight_layout()
    plt.grid(visible=1)
    plt.show()

#####################################################################