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

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]])

# Rayleigh range
def z_R(w_0:np.ndarray, lamb:float)->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)

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 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 single_gaussian_beam_potential_harmonic_approximation(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", depth:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um)->np.ndarray:
    U = - depth * (1 - (2 * (positions[0,:]/waists[0])**2) - (2 * (positions[2,:]/waists[1])**2) - (0.5 * positions[1,:]**2 * np.sum(np.reciprocal(z_R(waists, wavelength)))**2))
    return U
    
def harmonic_potential(pos, v, offset):
    U_Harmonic = ((0.5 * 164*u.u * (2 * np.pi * v*u.Hz)**2 * (pos*u.um)**2)/ac.k_B).to(u.uK) + offset*u.uK
    return U_Harmonic.value

def extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis):
    tmp_pos = Positions[axis, :]
    center_idx = np.where(tmp_pos == 0)[0][0]
    lb = int(round(center_idx - len(tmp_pos)/20, 1))
    ub = int(round(center_idx + len(tmp_pos)/20, 1))
    xdata = tmp_pos[lb:ub]
    tmp_pot = TrappingPotential[axis]
    Potential = tmp_pot[lb:ub]
    p0=[1e3, -TrapDepthInKelvin.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 plotHarmonicFit(Positions, TrappingPotential, TrapDepthInKelvin, 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([-TrapDepthInKelvin.value, max(TrappingPotential[axis].value)])
    plt.tight_layout()
    plt.grid(visible=1)
    plt.legend(prop={'size': 12, 'weight': 'bold'})
    plt.show()

def plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels):

    ## plot of the measured parameter vs. scan  parameter
    plt.figure(figsize=(9, 7))
    j = 0
    for i in range(np.size(ComputedPotentials, 0)):
        v, dv, popt, pcov = extractTrapFrequency(Positions, ComputedPotentials[i], TrapDepthInKelvin, axis)
        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])
        if i % 2 == 0 and j < len(Powers):
            plt.plot(Positions[axis], ComputedPotentials[i][axis], '--',label = 'P = ' + str(Powers[j]) + ' W; ' + TrapDepthLabels[j] + '; ' + tf_label) 
        elif i % 2 != 0 and j < len(Powers):
            plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[j]) + ' W; ' + tf_label) 
            j = j + 1
    if axis == 0:
        dir = 'X'
    elif axis == 1:
        dir = 'Y'
    else:
        dir = 'Z'
    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'})
    plt.show()
    # plt.savefig('pot_' + dir + '.png')

if __name__ == '__main__':

    # Powers = [0.1, 0.5, 2]
    # Powers = [5, 10, 20, 30, 40]
    Powers = [40]
    Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
    w_x, w_z = 34*u.um, 27.5*u.um # Beam Waists in the x and y directions
    # w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions
    # w_x, w_z = 20.5*u.um, 20.5*u.um

    axis = 1 # axis referenced to the beam along which you want the dipole trap potential
    extent = 1e4 # range of spatial coordinates in one direction to calculate trap potential over

    TrappingPotential = []
    ComputedPotentials = []
    TrapDepthLabels = []

    gravity = False
    astigmatism = True

    tilt_gravity = True
    theta = 1 # in degrees
    tilt_axis = [1, 0, 0] # lab space coordinates are rotated about x-axis in reference frame of beam

    disp_foci = 1.5 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um # difference in position of the foci along the propagation direction (Astigmatism)

    for p in Powers: 

        Power = p*u.W  # Single Beam Power
        
        TrapDepth = trap_depth(w_x, w_z, Power, alpha=Polarizability) 
        TrapDepthInKelvin = (TrapDepth/ac.k_B).to(u.uK)
        TrapDepthLabels.append("Trap Depth = " + str(round(TrapDepthInKelvin.value, 2)) + " " + str(TrapDepthInKelvin.unit))

        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.zeros((3, len(IdealTrappingPotential))) * IdealTrappingPotential.unit
        IdealTrappingPotential        = (IdealTrappingPotential/ac.k_B).to(u.uK) 

        if gravity and not astigmatism:
            ComputedPotentials.append(IdealTrappingPotential)
            # Influence of Gravity
            m = 164*u.u
            gravity_axis             = np.array([0, 0, 1]) 
            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) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
            
        elif not gravity and astigmatism:
            ComputedPotentials.append(IdealTrappingPotential)
            # Influence of Astigmatism
            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.zeros((3, len(TrappingPotential))) * TrappingPotential.unit
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        elif gravity and astigmatism:
            ComputedPotentials.append(IdealTrappingPotential)
            # Influence of Gravity and Astigmatism
            m = 164*u.u
            gravity_axis             = np.array([0, 0, 1]) 
            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) + gravitational_potential(gravity_axis_positions, m)
            TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
        
        else:
            TrappingPotential = IdealTrappingPotential

        # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis)
        # plotHarmonicFit(Positions, TrappingPotential, TrapDepthInKelvin, axis, popt, pcov)

        ComputedPotentials.append(TrappingPotential)
    
    # print(np.shape(ComputedPotentials))
    ComputedPotentials = np.asarray(ComputedPotentials)
    plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels)