from pylcp.atom import atom as base_atom
from pylcp.atom import state
import numpy as np
import scipy.constants as cts



class DyAtom(base_atom):
    def __init__(self, species):
        # Prepare to add in some useful electronic states:
        self.state = []

        if species == "161Dy" or species == "Dy161":
            self.I = 5/2  # nuclear spin
            self.gI = -0.1922  # nuclear magnetic moment
            self.mass = 160.93*cts.value('atomic mass constant')

            # Ground state:
            self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
                                    gJ=1.2415867, Ahfs=-116.2322e6,
                                    Bhfs =1091.5748e6, Chfs=0))
            #Blue Transition:
            self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
                                    tau=4.94e-9, gJ=1.22, Ahfs=-86.9e6,
                                    Bhfs=1747.4e6, Chfs=0))
            
        elif species == "163Dy" or species == "Dy163":
            self.I = 5/2  # nuclear spin
            self.gI = 0.269#-0.192  # nuclear magnetic moment
            self.mass = 162.93*cts.value('atomic mass constant')

            # Ground state:
            self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
                                    gJ=1.2415867, Ahfs=162.7543e6,
                                    Bhfs =1153.8684e6, Chfs=0))
            #Blue Transition:
            self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
                                    tau=4.94e-9, gJ=1.22, Ahfs=121.62e6,
                                    Bhfs=1844.9e6, Chfs=0))
            
        elif species == "164Dy" or species == "Dy164":
            self.I = 0  # nuclear spin
            self.gI = 0  # nuclear magnetic moment
            self.mass = 163.929*cts.value('atomic mass constant')

            # Ground state:
            self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
                                    gJ=1.25))
            #Blue Transition:
            self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
                                    tau=4.94e-9, gJ=1.22))
            
            #Red Transition:
            self.state.append(state(n=6, J=9, S=2, gJ=1.29, lam=626e-9, tau=1200e-9))


        else:
            # Fallback to normal atom behavior
            super().__init__(species)

        # Take the states and make transitions:
        base_atom._atom__make_transitions(self)