nx                  = 64; % Number of points along the x-axis
ny                  = 64; % Number of points along the y-axis
nz                  = 128; % Number of points along the z-axis

number_of_particles = 40000;

lx                  = 30;
ly                  = 20;
lz                  = 20;

x                   = linspace(-0.5*lx,0.5*lx-lx/nx,nx);
y                   = linspace(-0.5*ly,0.5*ly-ly/ny,ny);
z                   = linspace(-0.5*lz,0.5*lz-lz/nz,nz);

dx                  = x(end) - x(end-1);
dy                  = y(end) - y(end-1);
dz                  = z(end) - z(end-1);

psi                 = initialize_wavefunction(nx, ny, nz, x, y, z, number_of_particles, dx, dy, dz);

plot_wavefunction(psi, x, y, z);


%%
function [psi] = initialize_wavefunction(nx, ny, nz, x, y, z, number_of_particles, dx, dy, dz)
    
    psi = zeros(nx, ny, nz); % Initialize psi
    rng('shuffle'); % Seed random number generator
    
    for i = 1:nx
        for j = 1:ny
            for k = 1:nz
                random_number = rand();
                psi(i,j,k) = (1.0 + 0.1 * random_number) * ...
                    exp(-0.1 * (x(i)^2 + y(j)^2 + z(k)^2));
                
            end
        end
    end
    
    % Normalize psi
    norm_factor = sqrt(number_of_particles / (sum(abs(psi(:)).^2) * dx * dy * dz));
    psi = psi * norm_factor;
end

function plot_wavefunction(psi, x, y, z)
    % Function to plot the integrated wavefunction |psi|^2 along three planes
    
    % Compute the probability density
    psi_mag_sq = abs(psi).^2;

    % Integrate along each axis
    psi_xz = squeeze(trapz(y, psi_mag_sq, 2)); % Integrated over y
    psi_yz = squeeze(trapz(x, psi_mag_sq, 1)); % Integrated over x
    psi_xy = squeeze(trapz(z, psi_mag_sq, 3)); % Integrated over z
    
    %Plotting    
    figure(1);
    clf
    set(gcf,'Position', [100, 100, 1600, 450])
    t = tiledlayout(1, 3, 'TileSpacing', 'compact', 'Padding', 'compact'); % 1x3 grid

    nexttile;
    plotxz = pcolor(x,z,psi_xz');
    set(plotxz, 'EdgeColor', 'none');
    cbar1 = colorbar;
    cbar1.Label.Interpreter = 'latex';
    % cbar1.Ticks = []; % Disable the ticks
    colormap(gca, Helper.Colormaps.plasma())
    xlabel('$x$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    ylabel('$z$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    title('$|\Psi(x,z)|^2$', 'Interpreter', 'latex', 'FontSize', 14) 

    nexttile;
    plotyz = pcolor(y,z,psi_yz');
    set(plotyz, 'EdgeColor', 'none');
    cbar1 = colorbar;
    cbar1.Label.Interpreter = 'latex';
    % cbar1.Ticks = []; % Disable the ticks
    colormap(gca, Helper.Colormaps.plasma())
    xlabel('$y$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    ylabel('$z$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    title('$|\Psi(y,z)|^2$', 'Interpreter', 'latex', 'FontSize', 14) 

    nexttile;
    plotxy = pcolor(x,y,psi_xy');
    set(plotxy, 'EdgeColor', 'none');
    cbar1 = colorbar;
    cbar1.Label.Interpreter = 'latex';
    % cbar1.Ticks = []; % Disable the ticks
    colormap(gca, Helper.Colormaps.plasma())
    xlabel('$x$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    ylabel('$y$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
    title('$|\Psi(x,y)|^2$', 'Interpreter', 'latex', 'FontSize', 14) 

end
