import numpy as np
import scipy.constants as cts
from datetime import date

import pylcp
import common
import custom_rateeq
import Dy_atom
import PlotStuff
import pprint
import matplotlib.pyplot as plt

"""
This is an example how you could use the packege.
You can define a dictionary as one of these below and pass it on to the calc_force_field_and_trajectories_angled_2DMOT function
in common, if you want to calculate the force field and trajectories for a 2D MOT angled at 45° like in BoDy.
The results will be saved and plotted after the simulation is finished.
Change Isotope_name to use a different Isotope.
"""

Isotope_name = "Dy161"

params_DyLab = {
    "det" : -1.8,
    "s" : 1.76,
    "field_mag_Gauss_cm" : 34,
    "field_type" : "2DMOT",
    "waist_m" : 0.016,
    "aperture_m" : 0.0175,
    "trap_length_m" : 0.06,
    "max_velocity" : 150,
    "initial_velocities_m" : np.linspace(5,105,11),
    "initial_position_m" : -0.03,
    "boundary_m" : -0.03*1.1,
    "MOTvelocity_ms" : 1,
    "MOTradius_m": 0.001
}

params_BenLev = {
    "det" : -1.2,
    "s" : 0.185,
    "field_mag_Gauss_cm" : 10,
    "field_type" : "2DMOT",
    "waist_m" : 0.011,
    "aperture_m" : 0.01925,
    "trap_length_m" : 0.04,
    "max_velocity" : 150,
    "initial_velocities_m" : np.linspace(5,105,11),
    "initial_position_m" : -0.03,
    "boundary_m" : -0.03*1.1,
    "MOTvelocity_ms" : 0.1,
    "MOTradius_m": 0.0002
}

params_SingleBeamTest = {
    "det" : 0.,
    "s" : 0.,
    "polarisation" : 1,
    "field_mag_Gauss_cm" : 20,
    "field_type" : "Gradient",
    "waist_m" : 0.011,
    "aperture_m" : 0.01925,
    "trap_length_m" : 0.04,
    "max_velocity" : 150,
    "initial_velocities_m" : np.linspace(5,105,11),
    "initial_position_m" : -0.03,
    "boundary_m" : -0.03*1.1
}

#chose one of the above parameter sets:
params = params_DyLab

#define a memorable name for your results
name = "intoBoDy2DMOT"+str(params["det"])+"detuning"+str(params["s"])+"s"+str(params["field_mag_Gauss_cm"])+"Gauss_per_cm"

#do the calculation
common.calc_force_field_and_trajectories_angled_2DMOT(Isotope_name, params, resolution_x=200, resolution_v=200, t_max = 15e-3, basename=name)
#common.calc_force_field_and_trajectories_Single_Beam(Isotope_name, params_SingleBeamTest, 200, 200, 30e-3, name)

#And plot the results:
#PlotStuff.plot_force_field_and_trajectories(Isotope_name, name, date_str= "2025-04-28")
PlotStuff.plot_state_densities_and_force_field_with_trajectories(Isotope_name, name)#, date_str = "2025-04-28")
#PlotStuff.plot_force_field(Isotope_name, name, date_str= "2025-04-28")