%% Physical constants
PlanckConstant               = 6.62607015E-34;
PlanckConstantReduced        = 6.62607015E-34/(2*pi);
FineStructureConstant        = 7.2973525698E-3;
ElectronMass                 = 9.10938291E-31;
GravitationalConstant        = 6.67384E-11;
ProtonMass                   = 1.672621777E-27;
AtomicMassUnit               = 1.660539066E-27; 
BohrRadius                   = 5.2917721067E-11; 
BohrMagneton                 = 9.274009994E-24; 
BoltzmannConstant            = 1.38064852E-23;
StandardGravityAcceleration  = 9.80665;
SpeedOfLight                 = 299792458;
StefanBoltzmannConstant      = 5.670373E-8;
ElectronCharge               = 1.602176634E-19;
VacuumPermeability           = 1.25663706212E-6; 
DielectricConstant           = 8.8541878128E-12;
ElectronGyromagneticFactor   = -2.00231930436153;
AvogadroConstant             = 6.02214076E23;
ZeroKelvin                   = 273.15;
GravitationalAcceleration    = 9.80553;
VacuumPermittivity           = 1 / (SpeedOfLight^2 * VacuumPermeability);
HartreeEnergy                = ElectronCharge^2 / (4 * pi * VacuumPermittivity * BohrRadius);
AtomicUnitOfPolarizability   = (ElectronCharge^2 * BohrRadius^2) / HartreeEnergy; % Or simply 4*pi*VacuumPermittivity*BohrRadius^3

% Dy specific constants
Dy164Mass                    = 163.929174751*AtomicMassUnit;
Dy164IsotopicAbundance       = 0.2826;
DyMagneticMoment             = 9.93*BohrMagneton;

%% Example usage:
a_s          = 100 * BohrRadius; % Scattering length in meters
N            = 1e5;              % Number of atoms
v_x          = 10;
v_y          = 10;
v_z          = 1500;

[R_x, R_y, R_z] = calculateThomasFermiRadiusFromTrapParameters(v_x, v_y, v_z, N, a_s, Dy164Mass);
fprintf('Thomas-Fermi radii: R_x = %.2f µm, R_y = %.2f µm, R_z = %.2f µm\n', R_x*1e6, R_y*1e6, R_z*1e6);

%%

% Define the parameters
Power_values   = linspace(0.1, 1.5, 50);  % Laser powers (W), from 1 W to 10 W
waist_values   = linspace(50e-6, 90e-6, 50);  % Beam waists (m), from 10 µm to 100 µm

% Constants
a_s               = 100 * BohrRadius; % Scattering length in meters
N                 = 1e5;              % Number of atoms
LatticeConstant   = 13e-6;
alpha             = 180 * (AtomicUnitOfPolarizability / (2 * SpeedOfLight * VacuumPermittivity));
R_z_matrix        = zeros(length(waist_values), length(Power_values)); % Preallocate a matrix to store the R_z values
omega_eff_matrix  = zeros(length(waist_values), length(Power_values)); % Preallocate a matrix to store the omega_z values
% Loop over all combinations of P and w0
for i = 1:length(waist_values)
    for j = 1:length(Power_values)
        % Calculate the Thomas-Fermi radii for each P and w0
        [R_z, omega_eff] = calculateThomasFermiRadiusFromLatticeParameters(Power_values(j), waist_values(i), LatticeConstant, N, a_s, Dy164Mass, alpha);
        % Store the result in the matrix
        R_z_matrix(i, j)   = R_z * 1e6;  % Convert R_z to µm
        omega_eff_matrix(i, j)   = (omega_eff/(2*pi)) * 1e-3;  % Convert omega_z to kHz
    end
end

% Plot the resulting matrix
figure(1)
clf
set(gcf,'Position',[50 50 950 750])
set(gca,'FontSize',16,'Box','On','Linewidth',2);
imagesc(Power_values, waist_values * 1e6, R_z_matrix);  % Specify x and y data for axes
hold on
[contour_matrix, contour_handle] = contour(Power_values, waist_values * 1e6, R_z_matrix, 'LineColor', 'white', 'LineWidth', 1.5);
clabel(contour_matrix, contour_handle, 'FontSize', 12, 'Color', 'black', 'LabelSpacing', 600);  % Adjust LabelSpacing for better placement
set(gca, 'YDir', 'normal');              % Correct the y-axis direction
cbar = colorbar;                                % Add a colorbar
xlabel('Laser Power (W)','FontSize',16);
ylabel('Beam Waist (\mum)', 'Interpreter', 'tex','FontSize',16);
ylabel(cbar,'R_z (\mum)', 'Interpreter', 'tex','FontSize',16,'Rotation',270)
title(['\bf Thomas-Fermi Radius @ Lattice Spacing = ' num2str(LatticeConstant * 1E6) ' \mum'],'fontsize',16, 'Interpreter', 'tex');

% Plot the resulting matrix
figure(2)
clf
set(gcf,'Position',[50 50 950 750])
set(gca,'FontSize',16,'Box','On','Linewidth',2);
imagesc(Power_values, waist_values * 1e6, omega_eff_matrix);  % Specify x and y data for axes
hold on
[contour_matrix, contour_handle] = contour(Power_values, waist_values * 1e6, omega_eff_matrix, 'LineColor', 'white', 'LineWidth', 1.5);
clabel(contour_matrix, contour_handle, 'FontSize', 12, 'Color', 'black', 'LabelSpacing', 600);  % Adjust LabelSpacing for better placement
set(gca, 'YDir', 'normal');              % Correct the y-axis direction
cbar = colorbar;                                % Add a colorbar
xlabel('Laser Power (W)','FontSize',16);
ylabel('Beam Waist (\mum)', 'Interpreter', 'tex','FontSize',16);
ylabel(cbar,'\nu_z (kHz)', 'Interpreter', 'tex','FontSize',16,'Rotation',270)
title(['\bf Vertical Trap Frequency @ Lattice Spacing = ' num2str(LatticeConstant * 1E6) ' \mum'],'fontsize',16, 'Interpreter', 'tex');

%%
function [R_x, R_y, R_z] = calculateThomasFermiRadiusFromTrapParameters(v_x, v_y, v_z, N, a, m)
    % Calculate the Thomas-Fermi radius of a BEC.
    %
    % Parameters:
    % omega_x, omega_y, omega_z: Trap frequencies in x, y, z directions (Hz)
    % N: Number of atoms
    % a: s-wave scattering length (m)
    % m: Mass of the atom (kg)
    %
    % Returns:
    % R_x, R_y, R_z: Thomas-Fermi radii in meters

    omega_x = 2*pi*v_x;       % Trap frequency in x direction (Hz)
    omega_y = 2*pi*v_y;       % Trap frequency in y direction (Hz)
    omega_z = 2*pi*v_z;       % Trap frequency in z direction (Hz)

    % Calculate the geometric mean of trap frequencies
    omega_bar = (omega_x * omega_y * omega_z)^(1/3);

    % Constants
    PlanckConstantReduced        = 6.62607015E-34/(2*pi);                    % Reduced Planck's constant (J⋅s)

    % Calculate the mean harmonic oscillator length
    a_ho = sqrt(PlanckConstantReduced / (m * omega_bar));
    
    % Calculate the chemical potential
    mu = 0.5 * PlanckConstantReduced * omega_bar * (15 * N * a / a_ho)^(2/5);
    
    % Calculate Thomas-Fermi radii
    R_x = sqrt(2 * mu / (m * omega_x^2));
    R_y = sqrt(2 * mu / (m * omega_y^2));
    R_z = sqrt(2 * mu / (m * omega_z^2));
end

function [R_TF, omega_eff] = calculateThomasFermiRadiusFromLatticeParameters(P, w_z, d_lat, N, a, m, alpha)
    % Calculate Thomas-Fermi radius of a BEC in a Gaussian potential
    %
    % Inputs:
    % P     - Laser power (W)
    % w_z   - Beam waist (m)
    % theta - angle between beams
    % N - Number of atoms
    % a - s-wave scattering length (m)
    % m - Atomic mass (kg)
    %
    % Output:
    % R_TF - Thomas-Fermi radius (m)

    % Constants
    PlanckConstantReduced        = 6.62607015E-34/(2*pi);                     % Reduced Planck's constant (J⋅s)
    
    % Calculate trap depth
    I         = (2 * P) / (pi * w_z^2);
    U_single  = alpha * I;
    V_o       = 4 * U_single;
    
    % Calculate lattice constant
    % d_lat     = lambda / (2 * cos(deg2rad(theta)/2));
    
    % Calculate effective trapping frequency
    omega_eff = (pi / d_lat) * sqrt(2 * V_o / m);
    
    % Calculate chemical potential
    a_ho      = sqrt(PlanckConstantReduced / m * omega_eff);
    mu        = 0.5 * PlanckConstantReduced * omega_eff * (15 * N * a / a_ho)^(2/5); % Shallow lattice
    
    % Calculate Thomas-Fermi radius
    R_TF      = sqrt(2 * mu / (m * omega_eff^2));
end

