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

class ODP_1D:
    """
    1D optical dipole trap in prop_dir direction with gravity in z direction and
    magnetic field in mag_dir direction
    """
    def __init__(self, options):
        self.prop_dir = options["prop_dir"]
        self.wavelength = options["wavelength"]
        self.alpha = options["alpha"]*(4 * np.pi * ac.eps0 * ac.a0**3)/(2 * ac.eps0 * ac.c)  # in atomic units
        self.m = options["dy_mass"]
        self.mu = options["dy_mag_moment"]
        self.power = options["power"]  # laser power
        self.w_0 = options["w_0"]  # beam waist in transverse direction
        self.gravity = options["gravity"]  # boolean for gravity in z direction
        self.mag_dir = options["mag_dir"]  # direction of magnetic field and gradient
        self.mag_grad = options["mag_grad"]  # magnetic gradient
        self.mag_offset = options["mag_offset"]  # magnetic field offset
        self.mJ = options["mJ"] #magnetic quantum number
        self.ref_ind = options["ref_ind"] #refractive index

        #self.z_r = (np.pi * self.w_0**2 * self.ref_ind / self.wavelength).to(u.um) #rayleigh range
        #self.nu_r = (np.sqrt(2*self.power*self.alpha/np.pi/self.m) * 2/self.w_0**2).to(u.Hz)/(2*np.pi)
        #self.nu_z = (np.sqrt(4*self.power*self.alpha/np.pi**3/self.m) * self.wavelength/self.w_0**3/self.ref_ind).to(u.Hz)/(2*np.pi)

    @property
    def z_r(self):
        return (np.pi * self.w_0**2 * self.ref_ind / self.wavelength).to(u.um) #rayleigh range
    @property
    def nu_r(self):
        return (np.sqrt(2*self.power*self.alpha/np.pi/self.m) * 2/self.w_0**2).to(u.Hz)/(2*np.pi)
    @property
    def nu_z(self):
        return (np.sqrt(4*self.power*self.alpha/np.pi**3/self.m) * self.wavelength/self.w_0**3/self.ref_ind).to(u.Hz)/(2*np.pi)

    def get_axial_potential(self,z):
        """
        input:   z(array) positions to evaluate potential with unit of length
        output:  V(z)(array) potential energies at z
        """
        beam_pot = (-self.alpha * (2*self.power/np.pi/self.w_0**2) / (1+(z/self.z_r)**2)).to(u.J)
        #check if magnetic field is along axial direction
        if self.mag_dir == self.prop_dir:
            mag_grad_pot = (self.mu * self.mag_grad * z).to(u.J)
            mag_offset_pot = (self.mu * self.mag_offset * self.mJ).to(u.J)
        else:
            mag_grad_pot = 0
            mag_offset_pot = 0
        #check if gravity is on and axial direction is along z 
        if self.gravity and (self.prop_dir == "v"):
            gravity_pot = (self.m * ac.g0 * z).to(u.J)
        else:
            gravity_pot = 0

        return beam_pot + gravity_pot + mag_grad_pot + mag_offset_pot

    def get_radial_potential(self,r):
        """
        input:   r(array) positions to evaluate potential with unit of length
        output:  V(r)(array) potential energies at z
        """
        beam_pot = (-self.alpha * (2*self.power/np.pi/self.w_0**2) * np.exp(-2*r**2/self.w_0**2))
        #check if magnetic field is along radial direction
        if self.mag_dir != self.prop_dir:
            mag_grad_pot = (self.mu * self.mag_grad * r).to(u.J)
            mag_offset_pot = (self.mu * self.mag_offset * self.mJ).to(u.J)
        else:
            mag_grad_pot = 0
            mag_offset_pot = 0
        #check if gravity is on and axial direction is horizontal
        if self.gravity and (self.prop_dir == "h"):
            gravity_pot = (self.m * ac.g0 * r).to(u.J)
        else:
            gravity_pot = 0

        return beam_pot + gravity_pot + mag_grad_pot + mag_offset_pot
    
    def harm_pot(self, z, nu, V_0, z_0):
        return V_0 + 0.5*self.m.to(u.kg).value*(2*np.pi*nu)**2 * (z-z_0)**2
    
    def get_axial_occ_analytic(self,z,plot=False):
        """
        Finds the number of occupied energy states that are bound in the potential well,
        using the analytic harmonic approximation of the gaussian beam.
        """
        pot = self.get_axial_potential(z)
        if plot:
            plt.plot(z,(pot/ac.k_B).to(u.uK), label='actual potential')
        #find local minimum and maximum of potential
        try:
            E_max = pot[find_peaks(pot)[0]]
        except:
            raise Exception("potential maximum not found")
        try:
            ind_min = find_peaks(-pot)[0]
            E_min = pot[ind_min]
            z_min = z[ind_min][0]
        except:
            raise Exception("potential minimum not found")        
        
        delta_E = E_max - E_min
        #get eigenenergies of harmonic oscillator as temperatures
        n_max = int(np.ceil((delta_E/(ac.hbar*2*np.pi*self.nu_z)-0.5).value))
        n_temps = np.arange(n_max)
        temps = ((ac.hbar*2*np.pi*self.nu_z*(n_temps +0.5)+E_min)/ac.k_B).to(u.uK).value
        #temps = np.array([((ac.hbar*2*np.pi*self.nu_z*(n +0.5)+E_min)/ac.k_B).to(u.uK).value for n in range(n_max)])
        #plot results
        if plot:
            #plot energy levels
            plt.hlines(temps, np.min(z).to(u.um).value, np.max(z).to(u.um).value, linewidth=0.1,color="black", label="harmonic eigenenergies")
            #plot harmonic approximation
            #harm_approx = self.harm_pot(z.to(u.m).value, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J
            z_plot = np.array([zpt for zpt in z.to(u.m).value if self.harm_pot(zpt, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J < np.max(pot)])*u.m
            plt.plot(z_plot.to(u.um).value,(self.harm_pot(z_plot.to(u.m).value, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J/ac.k_B).to(u.uK),linestyle="--", label="harmonic approximation")
            plt.xlabel("z position [um]")
            plt.ylabel("trap depth [uK]")
            plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}, nu={self.nu_z}")
            plt.legend()
            plt.grid()
            plt.show()

        return n_max

    def get_axial_occ(self, z,fit_fac=1,p0=[1e8,-1.5e-27,0],plot=True,plot_fac=1):
        """
        #input:   z(array) positions
                  fit_fac(float) factor to make the interval for fitting the parabola bigger
                  p0(list) initial guess for the fit
                  plot(bool) enable/disable plotting
                  plot_fac(float) factor to make the plot interval for parabola bigger
        #output:  n_max(int) number of modes that fit into the trap
                  (nu(float) eigenfrequency)        
        """
        #get potential
        pot = self.get_axial_potential(z)
        if plot:
            plt.plot(z,(pot/ac.k_B).to(u.uK), label='actual potential')
        #get z values and potential in region to be fitted
        #z_min = z[find_peaks(-pot)[0]][0]
        #interval=(z_min - self.w_0, z_min + self.w_0)
        #z_fit = np.array([zpt for zpt in z.to(u.m).value if interval[0].to(u.m).value< zpt <interval[1].to(u.m).value])
        #z_fit = np.linspace(z_min - self.w_0, z_min + self.w_0,100).to(u.m).value
        z_fit = np.linspace((-self.w_0*fit_fac),(+self.w_0*fit_fac),100).to(u.m).value
        pot_fit = self.get_axial_potential(z_fit*u.m).to(u.J).value
        #perform fit
        try:
            popt, pcov = curve_fit(self.harm_pot, z_fit,pot_fit,p0=p0)
            nu = popt[0]*u.Hz
            print(nu)
            V_0 = popt[1]*u.J
            #get eigenenergies of harmonic oscillator as temperatures
            #find local maximum
            E_max = np.abs(pot[find_peaks(pot)[0]] - V_0)
            n_max = int(np.ceil((E_max/(ac.hbar*2*np.pi*nu)-0.5).value))
            temps = np.array([((ac.hbar*2*np.pi*nu*(n +0.5)+V_0)/ac.k_B).to(u.uK).value for n in range(n_max)])
            #print(f"{n_max} modes of trap frequncy {nu} were found in potential well.")
            #plot results
            if plot:
                plt.hlines(temps, np.min(z).to(u.um).value, np.max(z).to(u.um).value, linewidth=0.1,color="black", label="harmonic eigenenergies")
                z_plot = z_fit*u.m*plot_fac #np.linspace(interval[0],interval[1],100)
                plt.plot(z_plot.to(u.um).value,(self.harm_pot(z_plot.to(u.m).value, *popt)*u.J/ac.k_B).to(u.uK),linestyle="--", label="harmonic approximation")
                plt.xlabel("z position [um]")
                plt.ylabel("trap depth [uK]")
                plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}")
                plt.legend()
                plt.grid()
                plt.show()
        except:
            raise Exception("fit was unsuccessful :(")
            #plot potential for trouble shooting
            if plot:
                plt.xlabel("z position [um]")
                plt.ylabel("trap depth [uK]")
                plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}")
                plt.legend()
                plt.grid()
                plt.show()
                
        return n_max

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