import numpy as np
import matplotlib.pyplot as plt
import scipy.constants as cts
from pprint import pprint
import os
from datetime import date
import pylcp
import pickle
import custom_rateeq
import Dy_atom
from sympy import Rational
from pylcp.common import progressBar

"""
These are some functions common to all applications. These will be used for each run. Some of them are just for bookkeeping (checked_save() for example),
some of them are setting up the simulations for each run (init_hamiltonian...() for example) and some define specific simulation runs.
These will have to be extended for each case (Zeeman Slower / MOT / Mollasses etc.)
"""

def init_hamiltonian_singleF(F0, F1, gF0, gF1, mass, muB=1, gamma=1, k=1, returnBasis=False):
    """
    Create Hamiltonian for a two-level sytem of states F0 and F1

    Parameters
    ----------
    F0 : int or float
        Total angular momentum F of the ground state
    F1 : int or float
        Total angular momentum F of the excited state
    gF0 : float
        (F-)Landé g-factor of the ground state
    gF1 : float
        (F-)Landé g-factor of the excited state
    mass : float
        (unitless) Mass of the atom
    muB : float
        (unitless) Bohr magneton (default: 1)
    gamma : float
        (unitless) Natural linewidth of the excited state (default: 1)
    k : float
        (unitless) Wavenumber of the driving laser field (default: 1)
    returnBasis : bool, optional
        Whether to return the basis states used for the ground and excited
        state Hamiltonians (default: False)

    Returns
    -------
    Ham : pylcp.hamiltonian
        Hamiltonian for the 2 level system, ready to use in simulations
    E_g : np.ndarray
        Energies (eigenvalues) of the ground state
    E_e : np.ndarray
        Energies (eigenvalues) of the excited state
    basis_g : list (only if returnBasis is True)
        Basis states used for ground state Hamiltonian
    basis_e : list (only if returnBasis is True)
        Basis states used for excited state Hamiltonian
    """
    H_g, mu_q_g, basis_g = pylcp.hamiltonians.singleF(F=F0, gF=gF0, muB=muB, return_basis=returnBasis)
    H_e, mu_q_e, basis_e = pylcp.hamiltonians.singleF(F=F1, gF=gF1, muB=muB, return_basis=returnBasis)

    d_ijq = pylcp.hamiltonians.dqij_two_bare_hyperfine(F0, F1)

    Ham = pylcp.hamiltonian(H_g, H_e, mu_q_g, mu_q_e, d_ijq, mass = mass, muB=muB, gamma=gamma, k=k)

    E_g = np.unique(np.diagonal(H_g))
    E_e = np.unique(np.diagonal(H_e))

    if returnBasis:
        return Ham, E_g, E_e, basis_g, basis_e
    else:
        return Ham, E_g, E_e

def init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=1, gamma=1, k=1, state0=0, state1=1, returnBasis=False):
    """
    Create Hamiltonian for a two-level atom including hyperfine splitting

    Parameters
    ----------
    Isotope : object
        Atom object from PyLCP.atom class with nuclear magnetic moment, mass, states and transitions
    mass : float
        (unitless) Mass of the atom
    t0 : float, optional
        time unit, typically 1/Gamma
    muB : float, optional
        (unitless) Bohr magneton (default: 1)
    gamma : float, optional
        (unitless) Natural linewidth of the excited state (default: 1)
    k : float, optional
        (unitless) Wavenumber of the driving laser field (default: 1)
    state0 : int, optional
        Index of the ground state in Isotope.state (default: 0)
    state1 : int, optional
        Index of the excited state in Isotope.state (default: 1)
    returnStates : bool, optional
        Whether to return the basis states used in the Hamiltonians (default: False)

    Returns
    -------
    Hamiltonian : pylcp.hamiltonian
        Hamiltonian for the two-level hyperfine system
    E_g : np.ndarray
        Energies of the ground state manifold
    E_e : np.ndarray
        Energies of the excited state manifold
    basis_g : list (only if returnStates is True)
        Basis states used for ground state Hamiltonian
    basis_e : list (only if returnStates is True)
        Basis states used for excited state Hamiltonian
    """

    H_g, mu_q_g, basis_g = pylcp.hamiltonians.hyperfine_coupled(
        Isotope.state[state0].J, Isotope.I, Isotope.state[state0].gJ, Isotope.gI,
        Ahfs=Isotope.state[state0].Ahfs/Isotope.state[1].gammaHz,
        Bhfs=Isotope.state[state0].Bhfs/Isotope.state[1].gammaHz, Chfs=0, muB=muB, return_basis=True)
    
    H_e, mu_q_e, basis_e = pylcp.hamiltonians.hyperfine_coupled(
        Isotope.state[state1].J, Isotope.I, Isotope.state[state1].gJ, Isotope.gI,
        Ahfs=Isotope.state[state1].Ahfs/Isotope.state[1].gammaHz,
        Bhfs=Isotope.state[state1].Bhfs/Isotope.state[1].gammaHz, Chfs=0, muB=muB, return_basis= True)

    dijq = pylcp.hamiltonians.dqij_two_hyperfine_manifolds(
        Isotope.state[state0].J, Isotope.state[state1].J, Isotope.I)

    E_g = np.unique(np.diagonal(H_g))
    E_e = np.unique(np.diagonal(H_e))
    
    Hamiltonian = pylcp.hamiltonian(H_g, H_e, mu_q_g, mu_q_e, dijq, mass = mass, muB=muB, gamma=gamma, k=k)

    if returnBasis:
        return Hamiltonian, E_g, E_e, basis_g, basis_e

    else:
        return Hamiltonian, E_g, E_e
        
def init_LaserBeams2DMOT(detuning, beam_waist, aperture, k, saturation):
    laserBeams = pylcp.fields.laserBeams(
            [
                {
                    "kvec": k * np.array([+1, +1, 0.0])/np.sqrt(2),
                    "s": saturation,
                    "pol": -1,
                    "delta": detuning ,
                    "wb": beam_waist,
                    "rs": aperture,
                },
                {
                    "kvec": k * np.array([-1, -1, 0.])/np.sqrt(2),
                    "s": saturation,
                    "pol": -1,
                    "delta": detuning,
                    "wb": beam_waist,
                    "rs": aperture,
                },
                {
                    "kvec": k * np.array([1, -1, 0.0])/np.sqrt(2),
                    "s": saturation,
                    "pol": 1,
                    "delta": detuning,
                    "wb": beam_waist,
                    "rs": aperture,
                },
                {
                    "kvec": k * np.array([-1, 1,0.])/np.sqrt(2),
                    "s": saturation,
                    "pol": 1,
                    "delta": detuning,
                    "wb": beam_waist,
                    "rs": aperture,
                }
            ],
            beam_type=pylcp.fields.clippedGaussianBeam,
        )
    return laserBeams

def init_LaserBeamSingleZeeman(detuning, saturation, polarisation, beam_waist, aperture, k=1):
    laserBeam = pylcp.fields.laserBeams(
            [
                {
                    "kvec": k * np.array([-1., 0., 0.0]),
                    "s": saturation,
                    "pol": polarisation,
                    "delta": detuning ,
                    "wb": beam_waist,
                    "rs": aperture,
                }
            ],
            beam_type=pylcp.fields.clippedGaussianBeam,
        )
    return laserBeam

def init_1DMOTBeams(detuning, saturation, polarisation, beam_waist, aperture, k=1):
    laserBeam = pylcp.fields.laserBeams(
            [
                {
                    "kvec": k * np.array([-1., 0., 0.0]),
                    "s": saturation,
                    "pol": polarisation,
                    "delta": detuning ,
                    "wb": beam_waist,
                    "rs": aperture,
                },
                {
                    "kvec": k * np.array([1., 0., 0.0]),
                    "s": saturation,
                    "pol": polarisation,
                    "delta": detuning ,
                    "wb": beam_waist,
                    "rs": aperture,
                }
            ],
            beam_type=pylcp.fields.clippedGaussianBeam,
        )
    return laserBeam

def get_data_filepath(subfolder, base_name, extension):
    """
    Returns a unique file path inside the correct daily Data folder structure.
    If a file already exists under the same name an index will be appended to avoid overwriting files.
    
    Parameters:
    - subfolder: one of ['Plots', 'OtherData', 'Solutions']
    - base_name: the desired base file name (without extension)
    - extension: file extension including dot, e.g. '.npz'
    
    Returns:
    - full_path: a unique full path like Data/2025-03-17/OtherData/RateEqu_1.npz
    """
    assert subfolder in ['Plots', 'OtherData', 'Solutions'], "Invalid subfolder name!"

    # Create today's date string
    date_str = date.today().isoformat()

    # Build full folder path
    base_dir = os.path.join(".\Data", date_str)
    subfolder_path = os.path.join(base_dir, subfolder)

    # Create directories if needed
    os.makedirs(subfolder_path, exist_ok=True)

    # Generate filename
    filename = f"{base_name}{extension}"
    full_path = os.path.join(subfolder_path, filename)

    # Make filename unique
    counter = 1
    while os.path.exists(full_path):
        filename = f"{base_name}_{counter}{extension}"
        full_path = os.path.join(subfolder_path, filename)
        counter += 1

    return full_path

def get_data_filepath_and_check(subfolder, base_name, extension):
    """
    Returns a unique file path inside the correct daily Data folder structure.
    Checks if a file already excists and if yes waits for a user input.
    User can choose to replace existing or tp add an index to the path.
    
    Parameters:
    - subfolder: one of ['Plots', 'OtherData', 'Solutions']
    - base_name: the desired base file name (without extension)
    - extension: file extension including dot, e.g. '.npz'
    
    Returns:
    - full_path: a unique full path like Data/2025-03-17/OtherData/RateEqu_1.npz
    """
    assert subfolder in ['Plots', 'OtherData', 'Solutions'], "Invalid subfolder name!"

    # Create today's date string
    date_str = date.today().isoformat()

    # Build full folder path
    base_dir = os.path.join(".\Data", date_str)
    subfolder_path = os.path.join(base_dir, subfolder)

    # Create directories if needed
    os.makedirs(subfolder_path, exist_ok=True)

    # Generate filename
    filename = f"{base_name}{extension}"
    full_path = os.path.join(subfolder_path, filename)

    if os.path.exists(full_path):
        print(f"Warning: A file named '{filename}' already exists in the directory.")
        user_input = input("Do you want to replace the existing file? (yes/no): ").strip().lower()

        if user_input == 'yes':
            print("You chose to overwrite a file.")
            return True, full_path  # Return True to indicate overwrite permission
        else:
            # Make filename unique
            counter = 1
            while os.path.exists(full_path):
                filename = f"{base_name}_{counter}{extension}"
                full_path = os.path.join(subfolder_path, filename)
                counter += 1
            return False, full_path
    else:
        return False, full_path  # As the user did not explicitly agree to overwrite, return False.

def checked_save(check, path, data):
    if check:
        # Overwrite the file if check is True
        if path.endswith(".npz"):
            np.savez(path, **data)
            print(f"File '{os.path.basename(path)}' has been saved (overwrite allowed).")
        elif path.endswith(".pkl"):
            with open(path, 'wb') as f:
                pickle.dump(data, f)
            print(f"File '{os.path.basename(path)}' has been saved (overwrite allowed).")
        else:
            print("File type couldnt be recognised. Can only be '.pkl' or '.npz'.")
    else:
        # Ensure the file does not exist before saving
        if not os.path.exists(path):
            if path.endswith(".npz"):
                np.savez(path, **data)
                print(f"File saved as '{os.path.basename(path)}'.")
            elif path.endswith(".pkl"):
                with open(path, 'wb') as f:
                    pickle.dump(data, f)
                print(f"File saved as '{os.path.basename(path)}'.")
            else:
                print("File type couldnt be recognised. Can only be '.pkl' or '.npz'.")
        else:
            print(f"Error: Cannot save '{os.path.basename(path)}' because overwrite is not allowed.")

def calc_force_profile_Single_Beam(Isotope_name, params, resolution_x, resolution_v, basename):
   
    check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
    
    pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom

    Isotope = Dy_atom.DyAtom(Isotope_name)
    k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm 
    gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz

    # Define "natural units"
    x0 = 1/(k_lab) # cm
    t0 = 1/gamma_lab # s
    muB = 1
    k = k_lab*x0
    gamma = gamma_lab*t0
    mass = Isotope.mass*(x0/100)**2/(cts.h * t0)

    det = params["det"]
    s = params["s"]
    pol = params["polarisation"]
    field_mag = params["field_mag_Gauss_cm"]*1e-4
    field_type = params["field_type"]

    waist = params["waist_m"]/(x0/100)
    aperture = params["aperture_m"]/(x0/100)

    trap_length = params["trap_length_m"] #m
    zrange = trap_length/(x0/100)
    velocity = params["max_velocity"] # m/s
    vrange = velocity/((x0/100)/t0)
    v0s = params["initial_velocities_m"]/((x0/100)/t0)

    Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
    basis_g = np.transpose(basis_g)
    basis_e = np.transpose(basis_e)
    
    if Isotope_name == "Dy161":
        det_L = (E_e[0] - E_g[0]) + det
    elif Isotope_name == "Dy163":
        det_L = (E_e[-1] - E_g[-1]) + det
    elif Isotope_name == "Dy164":
        det_L = det
    else:
        raise(ValueError(f"Couldn't recognise Isotope name {Isotope_name}"))

    LaserBeams = init_LaserBeamSingleZeeman(detuning=det_L, saturation=s, polarisation=pol, beam_waist=waist, aperture=aperture, k=k)

    if field_type == "2DMOT":
        alpha_2DMOT = field_mag * cts.value('Bohr magneton') * t0 * x0  / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.magField(lambda R: [alpha_2DMOT*R[1],alpha_2DMOT*R[0],0])
    elif field_type == "Gradient":
        alpha_gradient = field_mag * cts.value('Bohr magneton') * t0 * x0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.magField(lambda R: [alpha_gradient*R[0],0,0])
    elif field_type == "Offset":
        alpha_offset = field_mag * cts.value('Bohr magneton') * t0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.constantMagneticField([alpha_offset,0,0])
    else:
        raise ValueError("Error: Unrecognised field type. Has to be either '2DMOT' or 'Gradient' or 'Offset'.")

    x = np.linspace(-zrange, zrange, resolution_x)
    v = np.linspace(-vrange, vrange, resolution_v)

    X, V = np.meshgrid(x, v)

    Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
    Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])

    print("Calculating rate equations:")
    Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
    print("Completed. \nCalculating Force Profile:")

    Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
    accel_profile = np.full((resolution_v, resolution_x), np.nan)

    progress = progressBar()
    count = 0
    num = resolution_v*resolution_x
    for i_v in range(resolution_v):
        for i_x in range(resolution_x):
            count += 1
            r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
            v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
            
            _, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
            magnitude = -1*np.sum(*R_ij['g->e'])/t0*cts.h*k_lab*100/Isotope.mass 
            accel_profile[i_v, i_x] = magnitude
            progress.update(count/num)

    data_dict = {
        'basis_g': basis_g,
        'basis_e': basis_e,
        'x0': x0,
        't0': t0,
        'x': x,
        'v': v,
        'acceleration_profile': accel_profile,
        'magField': field_mag
    }
    checked_save(check1, path1, data_dict)

def calc_force_field_and_trajectories_Single_Beam(Isotope_name, params, resolution_x, resolution_v, t_max, basename):
    
    check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
    
    Isotope = Dy_atom.DyAtom(Isotope_name)
    
    k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm 
    gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz

    # Define "natural units"
    x0 = 1/(k_lab) # cm
    t0 = 1/gamma_lab # s
    muB = 1
    k = k_lab*x0
    gamma = gamma_lab*t0
    mass = Isotope.mass*(x0/100)**2/(cts.h * t0)

    det = params["det"]
    s = params["s"]
    pol = params["polarisation"]
    field_mag = params["field_mag_Gauss_cm"]*1e-4
    field_type = params["field_type"]

    waist = params["waist_m"]/(x0/100)
    aperture = params["aperture_m"]/(x0/100)

    trap_length = params["trap_length_m"] #m
    zrange = trap_length/(x0/100)
    velocity = params["max_velocity"] # m/s
    vrange = velocity/((x0/100)/t0)
    v0s = params["initial_velocities_m"]/((x0/100)/t0)

    pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom

    Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
    basis_g = np.transpose(basis_g)
    basis_e = np.transpose(basis_e)
    
    if Isotope_name == "Dy164":
        det_L = det
    elif Isotope_name == "Dy163":
        det_L = (E_e[-1] - E_g[-1]) + det
    elif Isotope_name == "Dy161":
        det_L = (E_e[0] - E_g[0]) + det
    else:
        raise ValueError(f"Isotope name '{Isotope_name}' not recognised.")

    LaserBeams = init_LaserBeamSingleZeeman(detuning=det_L, saturation=s, polarisation=pol, beam_waist=waist, aperture=aperture, k=k)
    
    if field_type == "2DMOT":
        alpha_2DMOT = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.magField(lambda R: [-alpha_2DMOT*R[1],-alpha_2DMOT*R[0],0])
    elif field_type == "Gradient":
        alpha_gradient = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.magField(lambda R: [alpha_gradient*R[0],0,0])
    elif field_type == "Offset":
        alpha_offset = field_mag * cts.value('Bohr magneton') * t0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.constantMagneticField([alpha_offset,0,0])
    else:
        raise ValueError("Error: Unrecognised field type. Has to be either '2DMOT' or 'Gradient' or 'Offset'.")


    x = np.linspace(-zrange, zrange, resolution_x)
    v = np.linspace(-vrange, vrange, resolution_v)

    X, V = np.meshgrid(x, v)

    Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
    Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])

    print("Calculating rate equations:")
    Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
    print("Completed. \nCalculating Force Profile:")
    
    Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
    accel_profile = np.full((resolution_v, resolution_x), np.nan)

    progress = progressBar()
    count = 0
    num = resolution_v*resolution_x
    for i_v in range(resolution_v):
        for i_x in range(resolution_x):
            count += 1
            r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
            v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
            
            _, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
            magnitude = -1*np.sum(*R_ij['g->e'])/t0*cts.h*k_lab*100/Isotope.mass 
            accel_profile[i_v, i_x] = magnitude
            progress.update(count/num)

    data_dict = {
        'basis_g': basis_g,
        'basis_e': basis_e,
        'x0': x0,
        't0': t0,
        'x': x,
        'v': v,
        'acceleration_profile': accel_profile
    }

    checked_save(check1, path1, data_dict)

    initial_position = params["initial_position_m"]/(x0/100)#x[int(len(x)*(1-(aperture/zrange)))]
    boundary = params["boundary_m"]/(x0/100)#initial_position*1.2
    
    for i, v0 in enumerate(v0s):
        check2, path2 = get_data_filepath_and_check('Solutions', 'sol'+Isotope_name+basename+str(round(v0*(x0/(100*t0))))+'mps', '.pkl')

        Rateeq.set_initial_position_and_velocity(np.array([initial_position, 0., 0.]), np.array([v0, 0, 0]))
        if isinstance(Rateeq, pylcp.rateeq):
            Rateeq.set_initial_pop(np.ones(len(basis_e)+len(basis_g))/(len(basis_e)+len(basis_g)))

        print("Calculating Trajectory:", i+1, "out of ", len(v0s))
        Rateeq.evolve_motion([0., t_max/t0], max_step=1e-4/t0, boundary=boundary, progress_bar=True)
        sol = Rateeq.sol
        
        checked_save(check2, path2, sol)

def calc_force_field_and_trajectories_angled_2DMOT(Isotope_name, params, resolution_x, resolution_v, t_max, basename):
    
    check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
    
    pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom

    Isotope = Dy_atom.DyAtom(Isotope_name)
    k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm 
    gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz

    # Define "natural units"
    x0 = 1/(k_lab) # cm
    t0 = 1/gamma_lab # s
    muB = 1
    k = k_lab*x0
    gamma = gamma_lab*t0
    mass = Isotope.mass*(x0/100)**2/(cts.h * t0)

    det = params["det"]
    s = params["s"]
    field_mag = params["field_mag_Gauss_cm"]*1e-4
    field_type = params["field_type"]

    waist = params["waist_m"]/(x0/100)
    aperture = params["aperture_m"]/(x0/100)

    trap_length = params["trap_length_m"] #m
    zrange = trap_length/(x0/100)
    velocity = params["max_velocity"] # m/s
    vrange = velocity/((x0/100)/t0)
    v0s = params["initial_velocities_m"]/((x0/100)/t0)

    
    Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
    basis_g = np.transpose(basis_g)
    basis_e = np.transpose(basis_e)
    
    if Isotope_name == "Dy164":
        det_L = det
    elif Isotope_name == "Dy163":
        det_L = (E_e[-1] - E_g[-1]) + det
    elif Isotope_name == "Dy161":
        det_L = (E_e[0] - E_g[0]) + det
    else:
        raise ValueError(f"Isotope name '{Isotope_name}' not recognised.")

    LaserBeams = init_LaserBeams2DMOT(detuning=det_L, saturation=s, beam_waist=waist, aperture=aperture, k=k)
    
    if field_type == "2DMOT":
        alpha_2DMOT = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h  #field mag is given as Gauss/cm
        magField = pylcp.fields.magField(lambda R: [-alpha_2DMOT*R[1],-alpha_2DMOT*R[0],0])
    else:
        raise ValueError("Error: Unrecognised field type. Field Type has to be '2DMOT'.")

    x = np.linspace(-zrange, zrange, resolution_x)
    v = np.linspace(-vrange, vrange, resolution_v)

    X, V = np.meshgrid(x, v)

    Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
    Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])

    print("Calculating rate equations:")
    Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
    print("Completed. \nCalculating Force Profile:")
    
    Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
    accel_profile = np.full((resolution_v, resolution_x), np.nan)

    progress = progressBar()
    count = 0
    num = resolution_v*resolution_x
    for i_v in range(resolution_v):
        for i_x in range(resolution_x):
            count += 1
            r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
            v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
            
            _, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
            
            #Rij contains an entry for each laser beam. So we want to sum all contributtions of all beams into one array:
            summed_array = R_ij['g->e'].sum(axis=0)
            
            #And the magnitude of the force will be all contributions of all transiotions summed, times some factors.
            magnitude = -1*np.sum(summed_array)/t0*cts.h*k_lab*100/Isotope.mass 
            accel_profile[i_v, i_x] = magnitude
            progress.update(count/num)

    data_dict = {
        'basis_g': basis_g,
        'basis_e': basis_e,
        'x0': x0,
        't0': t0,
        'x': x,
        'v': v,
        'acceleration_profile': accel_profile
    }

    checked_save(check1, path1, data_dict)

    initial_position = params["initial_position_m"]/(x0/100)
    if "boundary_m" in params:
        boundary = params["boundary_m"]/(x0/100)
    else:
        boundary = None

    if "MOTvelocity_ms" in params:
        MOTvelocity = params["MOTvelocity_ms"]/((x0/100)/t0)
    else:
        MOTvelocity = None

    if "MOTradius_m" in params:
        MOTradius = params["MOTradius_m"]/(x0/100)
    else:
        MOTradius = None
    

    for i, v0 in enumerate(v0s):
        check2, path2 = get_data_filepath_and_check('Solutions', 'sol'+Isotope_name+basename+str(round(v0*(x0/(100*t0))))+'mps', '.pkl')

        Rateeq.set_initial_position_and_velocity(np.array([initial_position, 0., 0.]), np.array([v0, 0, 0]))
        if isinstance(Rateeq, pylcp.rateeq):
            Rateeq.set_initial_pop(np.ones(len(basis_e)+len(basis_g))/(len(basis_e)+len(basis_g)))

        print("Calculating Trajectory:", i+1, "out of ", len(v0s))
        Rateeq.evolve_motion([0., t_max/t0], max_step=1e-4/t0, boundary=boundary, MOTradius=MOTradius, MOTvelocity=MOTvelocity, progress_bar=True)
        sol = Rateeq.sol
        
        checked_save(check2, path2, sol)