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

#####################################################################
# 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 == 'arccos':
        phi         = np.linspace(0, 2*np.pi, n_points)
        first_half  = 2/np.pi * np.arccos(phi/np.pi-1) - 1
        second_half = np.flip(first_half)
        mod_func    = mod_amp * np.concatenate((first_half, second_half))
        dx          = (max(mod_func) - min(mod_func))/(2*n_points)
        return dx, mod_func
    else:
        return None

#####################################################################
# 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 meanThermalVelocity(T, m = 164*u.u):
    return 4 * np.sqrt((ac.k_B * T) /(np.pi * m))

def particleDensity(w_x, w_z, Power, Polarizability, N, T, m = 164*u.u): # For a thermal cloud
    v_x = calculateTrapFrequency(w_x, w_z, Power, Polarizability, dir = 'x')
    v_y = calculateTrapFrequency(w_x, w_z, Power, Polarizability, dir = 'y')
    v_z = calculateTrapFrequency(w_x, w_z, Power, Polarizability, 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 thermaldeBroglieWavelength(T, m = 164*u.u):
    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 = 9.93 * ac.muB, m = 164*u.u):
    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, Polarizability, N, T, B): #For a 3D Harmonic Trap
    return (particleDensity(w_x, w_z, Power, Polarizability, N, T) * scatteringCrossSection(B) * meanThermalVelocity(T) / (2 * np.sqrt(2))).decompose()

def calculatePSD(w_x, w_z, Power, Polarizability, N, T):
    return (particleDensity(w_x, w_z, Power, Polarizability, N, T, m = 164*u.u) * thermaldeBroglieWavelength(T)**3).decompose()

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

def gravitational_potential(positions: "np.ndarray|u.quantity.Quantity", m:"float|u.quantity.Quantity"):
    return m * ac.g0 * positions

def single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", alpha:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um)->np.ndarray:
    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: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", del_y:"float|u.quantity.Quantity", alpha:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um)->np.ndarray:
    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 harmonic_potential(pos, v, xoffset, yoffset, m = 164*u.u):
    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 modulated_single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", alpha:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um, mod_amp:"float|u.quantity.Quantity"=1)->np.ndarray:
    mod_amp = mod_amp * waists[0]
    n_points = int(len(positions[0,:])/2)
    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

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

def trap_depth(w_1:"float|u.quantity.Quantity", w_2:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity", alpha:float)->"float|u.quantity.Quantity":
    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, Polarizability, m = 164*u.u, dir = 'x'):
    TrapDepth = trap_depth(w_x, w_z, Power, alpha=Polarizability)
    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 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)/150, 1))
    ub = int(round(center_idx + len(tmp_pot)/150, 1))
    xdata = tmp_pos[lb:ub]
    Potential = tmp_pot[lb:ub]
    p0=[1e3, tmp_pos[center_idx].value, np.argmin(tmp_pot.value)]
    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, Polarizability, options):

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

    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, alpha=Polarizability) 
    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)
    
    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 
    Positions        = np.vstack((x_Positions, y_Positions, z_Positions)) * projection_axis[:, np.newaxis]

    IdealTrappingPotential        = single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power, alpha = Polarizability) 
    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 = 164*u.u
        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, alpha = Polarizability)
        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
        disp_foci = options['disp_foci']
        TrappingPotential        = astigmatic_single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power, del_y = disp_foci, alpha = Polarizability)
        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 = 164*u.u
        gravity_axis             = np.array([0, 0, 1]) 
        tilt_gravity             = options['tilt_gravity']
        theta                    = options['theta']
        tilt_axis                = options['tilt_axis']
        disp_foci                = options['disp_foci']
        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 = disp_foci, alpha = Polarizability)
        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
    
    if TrappingPotential[axis][0] > TrappingPotential[axis][-1]:
        EffectiveTrapDepthInKelvin = TrappingPotential[axis][-1] - min(TrappingPotential[axis])
    elif TrappingPotential[axis][0] < TrappingPotential[axis][-1]:
        EffectiveTrapDepthInKelvin = TrappingPotential[axis][0] - min(TrappingPotential[axis])
    else:
        EffectiveTrapDepthInKelvin = IdealTrapDepthInKelvin

    TrapDepthsInKelvin = [IdealTrapDepthInKelvin, EffectiveTrapDepthInKelvin]

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

    v, dv, popt, pcov = extractTrapFrequency(Positions, IdealTrappingPotential, axis)
    IdealTrapFrequency = [v, dv]
    v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, axis)
    TrapFrequency = [v, dv]
    ExtractedTrapFrequencies = [IdealTrapFrequency, TrapFrequency]

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

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

def plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, axis, popt, pcov):
    v = popt[0]
    dv = pcov[0][0]**0.5
    happrox = harmonic_potential(Positions[axis, :].value, *popt)
    plt.figure()
    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.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show()

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 plotPotential(Positions, ComputedPotentials, axis, Params = [], listToIterateOver = [], save = False):

    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]
        
        v = Params[j][2][1][0]
        dv = Params[j][2][1][1]
        
        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(Power, waists, alpha, wavelength, options):
        w_x = waists[0]
        w_z = waists[1]
        extent = options['extent']
        modulation = options['modulation']
        
        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]

            alpha = Polarizability
            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))
            I = A * np.exp(-2 * ((xm/w(ym, w_x, wavelength))**2 + (zm/w(ym, w_z, wavelength))**2))
            I = np.transpose(I.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)

            fig = plt.figure(figsize=(12, 6))
            ax = fig.add_subplot(121)
            ax.set_title('Intensity Profile ($MW/cm^2$)\n Aspect Ratio = %.2f' %(w_x/w_z))
            im = plt.imshow(I[:,idx,:].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(x_Positions==0)[0][0], idx, :], label = 'X - Horizontal')
            plt.plot(z_Positions, U[:, idx, np.where(z_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()    
        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  = int(len(x_Positions)/2)
            dx, xmod_Positions = modulation_function(mod_amp, n_points, func = 'arccos')
            
            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')
            A = 2*Power/(np.pi*w(0*u.um , w_x, wavelength)*w(0*u.um , 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 = np.transpose(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)

            fig = plt.figure(figsize=(12, 6))
            ax = fig.add_subplot(121)
            ax.set_title('Intensity Profile ($MW/cm^2$)')
            im = plt.imshow(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(x_Positions==0)[0][0], :], label = 'X - Horizontal')
            plt.plot(z_Positions, U[:, np.where(z_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 plotScatteringLengths():
    BField = np.arange(0, 2.59, 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()

#####################################################################
# RUN SCRIPT WITH OPTIONS BELOW                                     #
#####################################################################

if __name__ == '__main__':            

    Power = 40*u.W
    Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
    Wavelength = 1.064*u.um
    w_x, w_z = 27.5*u.um, 33.8*u.um # Beam Waists in the x and y directions

    # Power = 11*u.W
    # Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
    # w_x, w_z = 54.0*u.um, 54.0*u.um # Beam Waists in the x and y directions
    
    options = {
        'axis': 0, # axis referenced to the beam along which you want the dipole trap potential
        'extent': 3e2, # range of spatial coordinates in one direction to calculate trap potential over
        'modulation': True,
        'aspect_ratio': 3.67,
        'gravity': False,
        'tilt_gravity': False,
        'theta': 5, # in degrees
        'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam
        'astigmatism': False,
        'disp_foci': 3 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um # difference in position of the foci along the propagation direction (Astigmatism)
    }
    
    """Plot ideal and trap potential resulting for given parameters only"""
    # ComputedPotentials = [] 
    # Params = []  

    # Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Power, Polarizability, options)
    # ComputedPotentials.append(IdealTrappingPotential)
    # ComputedPotentials.append(TrappingPotential)
    # Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])

    # ComputedPotentials = np.asarray(ComputedPotentials)
    # plotPotential(Positions, ComputedPotentials, options['axis'], Params)

    """Plot harmonic fit for trap potential resulting for given parameters only"""
    # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, options['axis'])
    # plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, options['axis'], popt, pcov)

    """Plot trap potential resulting for given parameters (with one parameter being a list of values and the potential being computed for each of these values) only"""
    # ComputedPotentials = [] 
    # Params = []  
    # Power = [10, 20, 25, 30, 35, 40]*u.W # Single Beam Power
    # for p in Power: 
    #     Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, p, Polarizability, options)
    #     ComputedPotentials.append(IdealTrappingPotential)
    #     ComputedPotentials.append(TrappingPotential)
    #     Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])
        
    # ComputedPotentials = np.asarray(ComputedPotentials)
    # plotPotential(Positions, ComputedPotentials, options['axis'], Params)

    """Plot transverse intensity profile and trap potential resulting for given parameters only"""
    # options = {
    #     'extent': 70, # range of spatial coordinates in one direction to calculate trap potential over
    #     'modulation': True,
    #     'modulation_amplitude': 4.37
    # }

    # plotIntensityProfileAndPotentials(Power, [w_x, w_z], Polarizability, Wavelength, options)

    
    """Calculate relevant parameters for evaporative cooling"""
    AtomNumber = 1.13 * 1e7
    Temperature = 100 * u.uK
    BField = 2.1 * u.G
    modulation = False
    
    if modulation:
        aspect_ratio = 3.67
        init_ar = w_x / w_z
        w_x = w_x * (aspect_ratio / init_ar)

    n = particleDensity(w_x, w_z, Power, Polarizability, N = AtomNumber, T = Temperature, m = 164*u.u).decompose().to(u.cm**(-3))
    Gamma_elastic = calculateElasticCollisionRate(w_x, w_z, Power, Polarizability, N = AtomNumber, T = Temperature, B = BField)
    PSD = calculatePSD(w_x, w_z, Power, Polarizability, N = AtomNumber, T = Temperature).decompose()
    
    print('Particle Density = %.2E ' % (n.value) + str(n.unit))
    print('Elastic Collision Rate = %.2f ' % (Gamma_elastic.value) + str(Gamma_elastic.unit))
    print('PSD = %.2E ' % (PSD.value))

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

    print('v_x = %.2f ' %(v_x.value) + str(v_x.unit))
    print('v_y = %.2f ' %(v_y.value) + str(v_y.unit))
    print('v_z = %.2f ' %(v_z.value) + str(v_z.unit))

    print('a_s = %.2f ' %(scatteringLength(BField)[0] / ac.a0))
    
    # plotScatteringLengths()