%% 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;

%% 2-D DDI Potential in k-space, with Gaussian ansatz width determined by constrained minimization
% With user-defined values of interaction, density and tilt

% w0           = 2*pi*61.6316;                                                                  % Angular frequency unit [s^-1]
% l0           = sqrt(PlanckConstantReduced/(Dy164Mass*w0));
% % Defining a harmonic oscillator length - heredue to the choice of w0, l0
% is 1 micrometer

wz           = 2 * pi * 300;                                                                    % Trap frequency in the tight confinement direction
lz           = sqrt(PlanckConstantReduced/(Dy164Mass * wz));                                    % Defining a harmonic oscillator length

% Number of grid points in each direction
Params.Nx    = 256;
Params.Ny    = 256;
Params.Lx    = 150*1e-6;
Params.Ly    = 150*1e-6;
[Transf]     = setupSpace(Params);

nadd2s       = 0.110; % Number density * add^2
as_to_add    = 0.782; % 1/edd
Params.theta = 57;    % Polar angle of dipole moment
Params.phi   = 0;     % Azimuthal angle of dipole moment

add          = VacuumPermeability*DyMagneticMoment^2*Dy164Mass/(12*pi*PlanckConstantReduced^2); % Dipole length
gdd          = VacuumPermeability*DyMagneticMoment^2/3;

x0           = 5;
Aineq        = []; 
Bineq        = []; 
Aeq          = []; 
Beq          = [];
lb           = [1]; 
ub           = [10];
nonlcon      = [];
fminconopts  = optimoptions(@fmincon,'Display','off', 'StepTolerance', 1.0000e-11, 'MaxIterations',1500);

AtomNumberDensity      = nadd2s / add^2;                                                             % Number density of atoms
as                     = as_to_add * add;                                                            % Scattering length
eps_dd                 = add/as;                                                                     % Relative interaction strength
gs                     = 4 * pi * PlanckConstantReduced^2/Dy164Mass * as;                            % Contact interaction strength
TotalEnergyPerParticle = @(x) computeTotalEnergyPerParticle(x, as, AtomNumberDensity, wz, lz, gs, add, gdd, PlanckConstantReduced);
sigma                  = fmincon(TotalEnergyPerParticle, x0, Aineq, Bineq, Aeq, Beq, lb, ub, nonlcon, fminconopts);

MeanWidth         = sigma * lz;

% == 2-D DDI Potential in k-space == %
VDk               = compute2DPotentialInMomentumSpace(Transf, Params, MeanWidth); % Transf.K is in shifted space, so VDk is shifted
VDk_fftshifted    = fftshift(VDk);                                                % This is in normal space, but dispersion is calculated in shifted space, so unshifted VDk needs to be used there

figure(11)
clf
set(gcf,'Position',[50 50 950 750])
imagesc(fftshift(Transf.kx)*1e-6, fftshift(Transf.ky)*1e-6, VDk_fftshifted);      % Specify x and y data for axes
set(gca, 'YDir', 'normal');                                                       % Correct the y-axis direction
cbar1 = colorbar;
cbar1.Label.Interpreter = 'latex';
xlabel('$k_x l_o$','fontsize',16,'interpreter','latex');
ylabel('$k_y l_o$','fontsize',16,'interpreter','latex');
title(['2-D DDI Potential: $\theta = ',num2str(Params.alpha), '; \phi = ', num2str(Params.phi),'$'],'fontsize',16,'interpreter','latex')

% == Quantum Fluctuations term == %
gammaQF           = (32/3) * gs * (as^3/pi)^(1/2) * (1 + ((3/2) * eps_dd^2));
gamma5            = sqrt(2/5) / (sqrt(pi) * MeanWidth)^(3/2);
gQF               = gamma5 * gammaQF;

EpsilonK          = zeros(length(Transf.ky), length(Transf.kx));
gs_tilde          = gs / (sqrt(2*pi) * MeanWidth);
gdd_tilde         = gdd / (sqrt(2*pi) * MeanWidth); % you were missing this!

% == Dispersion relation == %
for idx = 1:length(Transf.kx)
    for jdx = 1:length(Transf.ky)
        DeltaK              = ((PlanckConstantReduced^2 .* (Transf.kx(idx).^2 + Transf.ky(jdx).^2)) ./ (2 * Dy164Mass)) + (2 * AtomNumberDensity * gs_tilde) + ((2 * AtomNumberDensity) * gdd_tilde .* VDk(jdx, idx)) + (3 * gQF * AtomNumberDensity^(3/2));
        EpsilonK(jdx, idx)  = sqrt(((PlanckConstantReduced^2 .* (Transf.kx(idx).^2 + Transf.ky(jdx).^2)) ./ (2 * Dy164Mass)) .* DeltaK); 
    end
end

EpsilonK = real(EpsilonK);

figure(12)
clf
set(gcf,'Position',[50 50 950 750])
imagesc(fftshift(Transf.kx)*1e-6, fftshift(Transf.ky)*1e-6, fftshift(EpsilonK));         % Specify x and y data for axes
set(gca, 'YDir', 'normal');                                                % Correct the y-axis direction
cbar1 = colorbar;
cbar1.Label.Interpreter = 'latex';
xlabel('$k_x l_o$','fontsize',16,'interpreter','latex');
ylabel('$k_y l_o$','fontsize',16,'interpreter','latex');
title(['2-D Dispersion: $\theta = ',num2str(Params.alpha), '; \phi = ', num2str(Params.phi),'$'],'fontsize',16,'interpreter','latex')

%% 3-D Plot
% Assuming Transf.kx, Transf.ky, and EpsilonK are already defined
% Apply fftshift and scaling
kx_shifted = fftshift(Transf.kx) * 1e-6;  % Convert to desired units
ky_shifted = fftshift(Transf.ky) * 1e-6;  % Convert to desired units
EpsilonK_shifted = fftshift(EpsilonK);    % This is the energy or function you want to plot

% Create meshgrid for the kx and ky values
[Kx, Ky] = meshgrid(kx_shifted, ky_shifted);

% Create the 3D plot
figure(13)
clf
set(gcf,'Position',[50 50 950 750])
surf(Kx, Ky, EpsilonK_shifted);
shading interp;  % Interpolates shading between grid points
% Label axes
xlabel('$k_x l_o$','fontsize',16,'interpreter','latex');
ylabel('$k_y l_o$','fontsize',16,'interpreter','latex');
zlabel('$\epsilon(k)$','fontsize',16,'interpreter','latex');
title(['2-D Dispersion: $\theta = ',num2str(Params.alpha), '; \phi = ', num2str(Params.phi),'$'],'fontsize',16,'interpreter','latex')
colorbar;  % Optional, to show a color scale for the surface plot
% Apply a different colormap
colormap('jet');  % 'jet', 'parula', 'hot', etc.
colorbar;  % Add colorbar to see the range of values

%% Cycle through angles

% Define values for theta and phi
theta_values      = 0:10:90;  % Range of theta values (you can modify this)
phi_values        = 0:10:90;   % Range of phi values (you can modify this)

% Set up VideoWriter object to produce a movie
% v = VideoWriter('potential_movie', 'MPEG-4'); % Create a video object
% v.FrameRate = 5;  % Frame rate of the video
% open(v);           % Open the video file

% Loop over theta and phi values
for theta = theta_values
    for phi = phi_values
        % Update Params with current theta and phi
        Params.theta = theta;
        Params.phi = phi;

        % Compute the potential for the current theta and phi
        % == 2-D DDI Potential in k-space == %
        VDk               = compute2DPotentialInMomentumSpace(Transf, Params, MeanWidth);
        VDk_fftshifted    = fftshift(VDk);
        
        % == Quantum Fluctuations term == %
        gammaQF           = (32/3) * gs * (as^3/pi)^(1/2) * (1 + ((3/2) * eps_dd^2));
        gamma5            = sqrt(2/5) / (sqrt(pi) * MeanWidth)^(3/2);
        gQF               = gamma5 * gammaQF;
        
        EpsilonK          = zeros(length(Transf.ky), length(Transf.kx));
        gs_tilde          = gs / (sqrt(2*pi) * MeanWidth);
        gdd_tilde         = gdd / (sqrt(2*pi) * MeanWidth); % you were missing this!
        
        % == Dispersion relation == %
        for idx = 1:length(Transf.kx)
            for jdx = 1:length(Transf.ky)
                DeltaK              = ((PlanckConstantReduced^2 .* (Transf.kx(idx).^2 + Transf.ky(jdx).^2)) ./ (2 * Dy164Mass)) + (2 * AtomNumberDensity * gs_tilde) + ((2 * AtomNumberDensity) * gdd_tilde .* VDk(jdx, idx)) + (3 * gQF * AtomNumberDensity^(3/2));
                EpsilonK(jdx, idx)  = sqrt(((PlanckConstantReduced^2 .* (Transf.kx(idx).^2 + Transf.ky(jdx).^2)) ./ (2 * Dy164Mass)) .* DeltaK); 
            end
        end

        EpsilonK = real(EpsilonK);

        % Plot the result
        figure(14)
        clf
        set(gcf,'Position',[50 50 950 750])
        imagesc(fftshift(Transf.kx)*1e-6, fftshift(Transf.ky)*1e-6, fftshift(EpsilonK));   % Specify x and y data for axes
        set(gca, 'YDir', 'normal');              % Correct the y-axis direction
        cbar1 = colorbar;
        cbar1.Label.Interpreter = 'latex';
        xlabel('$k_x l_o$','fontsize',16,'interpreter','latex');
        ylabel('$k_y l_o$','fontsize',16,'interpreter','latex');
        title(['2-D Dispersion: $\theta = ',num2str(Params.theta), '; \phi = ', num2str(Params.phi),'$'],'fontsize',16,'interpreter','latex')

        % Capture the frame and write to video
        % frame = getframe(gcf); % Capture the current figure
        % writeVideo(v, frame);  % Write the frame to the video
    end
end

% Close the video file
% close(v);

%%
function [Transf] = setupSpace(Params)
    
    Transf.Xmax            = 0.5*Params.Lx;
    Transf.Ymax            = 0.5*Params.Ly;
        
    Nx                     = Params.Nx; 
    Ny                     = Params.Ny;
    
    % Fourier grids
    x                      = linspace(-0.5*Params.Lx,0.5*Params.Lx-Params.Lx/Params.Nx,Params.Nx);
    Kmax                   = pi*Params.Nx/Params.Lx;
    kx                     = linspace(-Kmax,Kmax,Nx+1);
    kx                     = kx(1:end-1); 
    dkx                    = kx(2)-kx(1);
    kx                     = fftshift(kx);
    
    y                      = linspace(-0.5*Params.Ly,0.5*Params.Ly-Params.Ly/Params.Ny,Params.Ny);
    Kmax                   = pi*Params.Ny/Params.Ly;
    ky                     = linspace(-Kmax,Kmax,Ny+1);
    ky                     = ky(1:end-1); 
    dky                    = ky(2)-ky(1);
    ky                     = fftshift(ky);
    
    [Transf.X,Transf.Y]    = ndgrid(x,y);
    [Transf.KX,Transf.KY]  = ndgrid(kx,ky);
    Transf.x               = x; 
    Transf.y               = y;
    Transf.kx              = kx;
    Transf.ky              = ky;
    Transf.dx              = x(2)-x(1); 
    Transf.dy              = y(2)-y(1);
    Transf.dkx             = dkx; 
    Transf.dky             = dky;
end

function VDk = compute2DPotentialInMomentumSpace(Transf, Params, MeanWidth)
% == Calculating the DDI potential in Fourier space with appropriate cutoff == %
    % Angles of the dipole moment are defined in and away from the X-Z plane
    % Interaction in K space
    QX              = Transf.KX*MeanWidth/sqrt(2);
    QY              = Transf.KY*MeanWidth/sqrt(2);

    Qsq             = QX.^2 + QY.^2;
    absQ            = sqrt(Qsq);
    QDsq            = QX.^2*cos(Params.phi)^2 + QY.^2*sin(Params.phi)^2; % phi here is the azimuthal angle of the dipole moment (angle from the x-axis)
    
    % Bare interaction
    Fpar            = -1 + 3*sqrt(pi)*QDsq.*erfcx(absQ)./absQ; % Scaled complementary error function erfcx(x) = e^(x^2) * erfc(x)
    Fperp           = 2 - 3*sqrt(pi).*absQ.*erfcx(absQ);
    Fpar(absQ == 0) = -1; 
    
    % Full DDI
    VDk             = (Fpar*sin(Params.theta)^2 + Fperp*cos(Params.theta)^2); % theta here is the polar angle of the dipole moment (angle from the z-axis)
end

function ret = computeTotalEnergyPerParticle(x, as, AtomNumberDensity, wz, lz, gs, add, gdd, PlanckConstantReduced)
    eps_dd                     = add/as;                                                                     % Relative interaction strength
    MeanWidth                  = x * lz;
    gammaQF                    = (32/3) * gs * (as^3/pi)^(1/2) * (1 + ((3/2) * eps_dd^2));                   % Quantum Fluctuations term
    gamma4                     = 1/(sqrt(2*pi) * MeanWidth);
    gamma5                     = sqrt(2/5) / (sqrt(pi) * MeanWidth)^(3/2);
    gQF                        = gamma5 * gammaQF;
    Energy_AxialComponent      = (PlanckConstantReduced * wz) * ((lz^2/(4 * MeanWidth^2)) + (MeanWidth^2/(4 * lz^2)));
    Energy_TransverseComponent = (0.5 * (gs + (2*gdd)) * gamma4 * AtomNumberDensity) + ((2/5) * gQF * AtomNumberDensity^(3/2));
    ret                        = (Energy_AxialComponent + Energy_TransverseComponent) / (PlanckConstantReduced * wz);
end



