%% Parameters

groupList     = ["/images/MOT_3D_Camera/in_situ_absorption", "/images/ODT_1_Axis_Camera/in_situ_absorption", "/images/ODT_2_Axis_Camera/in_situ_absorption", "/images/Horizontal_Axis_Camera/in_situ_absorption", "/images/Vertical_Axis_Camera/in_situ_absorption"];

folderPath    = "C:/Users/Karthik/Documents/GitRepositories/Calculations/Imaging-Response-Function-Extractor/";

run           = '0072';

folderPath    = strcat(folderPath, run);

cam           = 5; 

angle         = 0;
center        = [1141, 2049];
span          = [255, 255];
fraction      = [0.1, 0.1];

NA            = 0.6;
Magnification = 21;
pixel_size    = 4.6e-6 / Magnification;
lambda        = 421e-9;

% The diameter of the first dark concentric ring surrounding the central intensity peak of a point spread function (or Airy disk)
d             = 1.22 * (lambda / NA); 

% Abbe limit
AbbeLimit     = lambda / (2 * NA); 

% Maximum resolvable spatial frequency for the coherent case
k_cutoff      = 2 * pi * (NA/lambda) * 1e-6; % (in units of 1/µm)

removeFringes = false;

%% Compute OD image, rotate and extract ROI for analysis
% Get a list of all files in the folder with the desired file name pattern.

filePattern = fullfile(folderPath, '*.h5');
files       = dir(filePattern);
refimages   = zeros(span(1) + 1, span(2) + 1, length(files));
absimages   = zeros(span(1) + 1, span(2) + 1, length(files));

for k = 1 : length(files)
    baseFileName = files(k).name;
    fullFileName = fullfile(files(k).folder, baseFileName);
    
    fprintf(1, 'Now reading %s\n', fullFileName);
    
    atm_img  = im2double(imrotate(h5read(fullFileName, append(groupList(cam), "/atoms")), angle));
    bkg_img  = im2double(imrotate(h5read(fullFileName, append(groupList(cam), "/background")), angle));
    dark_img = im2double(imrotate(h5read(fullFileName, append(groupList(cam), "/dark")), angle));
    
    refimages(:,:,k)  = subtractBackgroundOffset(cropODImage(bkg_img, center, span), fraction)';
    absimages(:,:,k)  = subtractBackgroundOffset(cropODImage(calculateODImage(atm_img, bkg_img, dark_img), center, span), fraction)';

end

% Fringe removal

if removeFringes
    optrefimages               = removefringesInImage(absimages, refimages);
    absimages_fringe_removed   = absimages(:, :, :) - optrefimages(:, :, :);
    
    nimgs                      = size(absimages_fringe_removed,3);
    od_imgs                    = cell(1, nimgs);
    for i = 1:nimgs
        od_imgs{i}             = absimages_fringe_removed(:, :, i);
    end
else
    nimgs                      = size(absimages(:, :, :),3);
    od_imgs                    = cell(1, nimgs);
    for i = 1:nimgs
        od_imgs{i}             = absimages(:, :, i);
    end
end

%% Compute the Density Noise Spectrum

mean_subtracted_od_imgs = cell(1, length(od_imgs));
mean_od_img             = mean(cat(3, od_imgs{:}), 3, 'double');

density_fft             = cell(1, length(od_imgs));
density_noise_spectrum  = cell(1, length(od_imgs));

[Nx, Ny] = size(mean_od_img);
dx       = pixel_size;
dy       = pixel_size;

xvals    = (1:Nx)*dx*1e6;
yvals    = (1:Ny)*dy*1e6;

dvx      = 1 / (Nx * dx);  % reciprocal space increment in the X direction (in units of 1/dx)
dvy      = 1 / (Ny * dy);  % reciprocal space increment in the Y direction (in units of 1/dy)

% Create the frequency axes
vx       = (-Nx/2:Nx/2-1) * dvx;  % Frequency axis in X (in units of 1/dx)
vy       = (-Ny/2:Ny/2-1) * dvy;  % Frequency axis in Y (in units of 1/dy)

% Create the Wavenumber axes
kx       = 2*pi*vx;  % Wavenumber axis in X
ky       = 2*pi*vy;  % Wavenumber axis in Y

% Create Circular Mask
n        = 2^8;                    % size of mask
fftmask  = zeros(n);
I        = 1:n; 
x        = I-n/2;                  % mask x-coordinates 
y        = n/2-I;                  % mask y-coordinates
[X,Y]    = meshgrid(x,y);          % create 2-D mask grid
R        = 32;                     % aperture radius
A          = (X.^2 + Y.^2 <= R^2);   % circular aperture of radius R
fftmask(A) = 1;                      % set mask elements inside aperture to 1

% Calculate Power Spectrum and plot
figure(1)
clf
set(gcf,'Position',[100, 100, 1200, 800])

% Create tiled layout with 2 rows and 3 columns
t = tiledlayout(2, 3, 'TileSpacing', 'compact', 'Padding', 'compact');

for k = 1 : length(od_imgs)
    mean_subtracted_od_imgs{k} = od_imgs{k} - mean_od_img;
    masked_img                 = mean_subtracted_od_imgs{k} .* fftmask;
    density_fft{k}             = (1/numel(masked_img)) * abs(fftshift(fft2(masked_img)));
    density_noise_spectrum{k}  = density_fft{k}.^2;
    
    % Tile 1: Single-shot image
    nexttile(1);
    imagesc(xvals, yvals, od_imgs{k})
    xlabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('Single-shot image', 'FontSize', 16);
    
    % Tile 2: Averaged density image
    nexttile(2);
    imagesc(xvals, yvals, mean_od_img)
    xlabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('Averaged density image', 'FontSize', 16);
    
    % Tile 3: Image noise = Single-shot - Average
    nexttile(3);
    imagesc(xvals, yvals, mean_subtracted_od_imgs{k})
    xlabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('Image noise = Single-shot - Average', 'FontSize', 16);
    
    % Tile 4: Masked Noise
    nexttile(4);
    imagesc(xvals, yvals, mean_subtracted_od_imgs{k} .* fftmask)
    xlabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('\mum', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('Masked Noise', 'FontSize', 16);
    
    % Tile 5: DFT
    nexttile(5);
    imagesc(kx*1E-6, ky*1E-6, abs(log2(density_fft{k})))
    xlabel('k_x (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('k_y (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('DFT', 'FontSize', 16);
    
    % Tile 6: Density Noise Spectrum = |DFT|^2
    nexttile(6);
    imagesc(kx*1E-6, ky*1E-6, abs(log2(density_noise_spectrum{k})))
    xlabel('k_x (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('k_y (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap jet;
    set(gca,'YDir','normal')
    title('Density Noise Spectrum', 'FontSize', 16);
    
    drawnow;
end

%% Compute the average 2D spectrum

% Compute the average power spectrum.
averagePowerSpectrum = mean(cat(3, density_noise_spectrum{:}), 3, 'double');

% Plot the average power spectrum and the 1-D Radial Imaging Response Function
figure(2)
clf
set(gcf,'Position',[100, 100, 1500, 700])

% Create tiled layout with 2 rows and 3 columns
t = tiledlayout(1, 2, 'TileSpacing', 'compact', 'Padding', 'compact');

nexttile(1);
imagesc(abs(10*log10(averagePowerSpectrum)))
axis equal tight;
colorbar
colormap(flip(jet));
title('Average Density Noise Spectrum', 'FontSize', 16);
grid on;

centers      = ginput;
radius       = 3;
% Plot where clicked
hVC          = viscircles(centers, radius, 'Color', 'r', 'LineWidth', 2);
xc           = centers(:,1);
yc           = centers(:,2);
[yDim, xDim] = size(averagePowerSpectrum);
[xx,yy]      = meshgrid(1:yDim,1:xDim);
MaskToRemoveUnwantedPeaks         = false(xDim,yDim);
for ii = 1:size(centers, 1)
	MaskToRemoveUnwantedPeaks     = MaskToRemoveUnwantedPeaks | hypot(xx - xc(ii), yy - yc(ii)) <= radius;
end
MaskToRemoveUnwantedPeaks         = not(MaskToRemoveUnwantedPeaks);

% Ask user if the circle is acceptable.
message = sprintf('Is this acceptable?');
button = questdlg(message, message, 'Accept', 'Reject and Quit', 'Accept');
if contains(button, 'Accept','IgnoreCase',true)
    image = MaskToRemoveUnwantedPeaks.*averagePowerSpectrum;
    image(image==0) = NaN;
    imagesc(kx*1E-6, ky*1E-6, MaskToRemoveUnwantedPeaks.*abs(10*log10(averagePowerSpectrum)))
    hold on
    xlabel('k_x (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('k_y (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap(flip(jet));
    title('Average Density Noise Spectrum', 'FontSize', 16);
    grid on;
elseif contains(button, 'Quit','IgnoreCase',true)
    delete(hVC); % Delete the circle from the overlay.
    image = averagePowerSpectrum;
    imagesc(kx*1E-6, ky*1E-6, abs(10*log10(averagePowerSpectrum)))
    xlabel('k_x (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    ylabel('k_y (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
    axis equal tight;
    colorbar
    colormap(flip(jet));
    title('Average Density Noise Spectrum', 'FontSize', 16);
    grid on;
end

% Fit
nexttile(2);
radialStep         = 1; % in pixels
[R, Profile]       = getRadialProfile(averagePowerSpectrum, radialStep);
kvec               = (2 * pi) .* R(2:end) .* (sqrt(dvx^2 + dvy^2) * radialStep) * 1E-6; % in units of micrometers^-1
NormalisedProfile  = (Profile(2:end) - min(Profile(2:end)))./(max(Profile(2:end)) - min(Profile(2:end)));
kmax               = k_cutoff;

% Define the objective function to minimize (the difference between model and data)
objectiveFunction  = @(C, k) RadialImagingResponseFunction(C, k, kmax) - NormalisedProfile;

% Initial guess for the parameters [C1, C2, C3, C4, C5, C6] 
initialGuess       = [2E6, 1E-6, 1E-6, 1E-6, 1E-6, 0.8];

% Set upper and lower bounds for the parameters (optional)
lb                 = [-Inf, -Inf, -Inf, -Inf, -Inf, 0];  % Lower bounds
ub                 = [Inf, Inf, Inf, Inf, Inf, 1];       % Upper bounds

% Perform the non-linear least squares fitting using lsqcurvefit
options            = optimoptions('lsqcurvefit', 'Display', 'iter');  % Display iterations during fitting
[C_fit, resnorm]   = lsqcurvefit(objectiveFunction, initialGuess, kvec, NormalisedProfile, lb, ub, options);

% Plot the fitted result against the data
k_new              = linspace(kvec(1), kvec(end), 1000);
fittedProfile      = RadialImagingResponseFunction(C_fit, k_new, kmax);

nexttile(2)
plot(kvec, NormalisedProfile, 'o-', 'MarkerSize', 4, 'MarkerFaceColor', 'none', 'DisplayName', 'Radial (average) profile');
hold on
plot(k_new, fittedProfile, 'r-', 'DisplayName', 'Fit');
set(gca, 'XScale', 'log'); % Setting x-axis to log scale
xlabel('k_\rho (\mum^{-1})', 'Interpreter', 'tex', 'FontSize', 16)
ylabel('Normalised amplitude', 'FontSize', 16)
title('Imaging Response Function', 'FontSize', 16);
legend('FontSize', 16);
grid on;

%% 2-D Imaging Response Function fit

% Get the size of the matrix
[nRows, nCols] = size(averagePowerSpectrum);

% Compute the center index
centerRow = round(nRows / 2);
centerCol = round(nCols / 2);

% Remove the central DC component by setting it to zero
averagePowerSpectrum(centerRow, centerCol) = 0;

% Define the objective function for fitting
function residuals = fitImagingResponse(params, averagePowerSpectrum)
    % Calculate the ImagingResponseFunction with the given parameters
    M = ImagingResponseFunction(params);

    % Compute the residuals as the difference between the matrices
    residuals = M(:) - averagePowerSpectrum(:); % Flatten both matrices
end

% Initial guess for parameters (from your example)
initialParams = [0.65, 0.1, 0.1, 0.1, 0.1, 0.1, 1E-6, 1E-8, 80];

% Set options for the optimizer (optional, helps with debugging)
options       = optimoptions('lsqnonlin', 'Display', 'iter', 'MaxFunctionEvaluations', 5000);

% Perform the fitting using lsqnonlin
optimalParams                  = lsqnonlin(@(params) fitImagingResponse(params, averagePowerSpectrum), initialParams, [], [], options);
AstigmatismCoefficient         = optimalParams(2);
SphericalAberrationCoefficient = optimalParams(3);
DefocussingCoefficient         = optimalParams(5);

% Compute the best-fit M using the optimized parameters
bestFitM  = ImagingResponseFunction(optimalParams);

residuals = fitImagingResponse(optimalParams, averagePowerSpectrum);

% Plot the fitted result

figure(3)
clf
set(gcf,'Position',[100, 100, 1500, 500])

% Create tiled layout with 2 rows and 3 columns
t = tiledlayout(1, 3, 'TileSpacing', 'compact', 'Padding', 'compact');

nexttile(1);
imagesc(abs(10*log10(averagePowerSpectrum)))
axis equal tight;
colorbar
colormap(flip(jet));
title('Average Density Noise Spectrum', 'FontSize', 16);
grid on;

nexttile(2);
imagesc(abs(10*log10(bestFitM)))
axis equal tight;
colorbar
colormap(flip(jet));
title('Fitted Imaging Response Function', 'FontSize', 16);
grid on;

nexttile(3);
resmat = reshape(residuals, size(bestFitM));
imagesc(abs(10*log10(resmat)))
axis equal tight;
colorbar
colormap(flip(jet));
title('Fit residuals', 'FontSize', 16);
grid on;

%% 3-D Plot of Density Noise Spectrum

figure(4)
clf
set(gcf,'Position',[100, 100, 1500, 700])
surf(10 * log10(averagePowerSpectrum));
shading interp;   % Creates a 3D surface plot
xlabel('X-axis', 'FontSize', 16);
ylabel('Y-axis', 'FontSize', 16);
zlabel('Rescaled amplitude (a.u.)', 'FontSize', 16);
title('Imaging Response Function', 'FontSize', 16);
colorbar;         % Shows the color scale

% Set axis limits based on the size of the averagePowerSpectrum
[xSize, ySize] = size(averagePowerSpectrum);
xlim([1, xSize]);
ylim([1, ySize]);
zlim([min(10 * log10(averagePowerSpectrum(:))), max(10 * log10(averagePowerSpectrum(:)))]);  % Optional for Z-axis limits

%% Helper Functions

function ret = getBkgOffsetFromCorners(img, x_fraction, y_fraction)
    % image must be a 2D numerical array
    [dim1, dim2] = size(img);

    s1 = img(1:round(dim1 * y_fraction), 1:round(dim2 * x_fraction)); 
    s2 = img(1:round(dim1 * y_fraction), round(dim2 - dim2 * x_fraction):dim2);
    s3 = img(round(dim1 - dim1 * y_fraction):dim1, 1:round(dim2 * x_fraction));
    s4 = img(round(dim1 - dim1 * y_fraction):dim1, round(dim2 - dim2 * x_fraction):dim2);

    ret = mean([mean(s1(:)), mean(s2(:)), mean(s3(:)), mean(s4(:))]);
end

function ret = subtractBackgroundOffset(img, fraction)
    % Remove the background from the image.
    % :param dataArray: The image
    % :type dataArray: xarray DataArray
    % :param x_fraction: The fraction of the pixels used in x axis
    % :type x_fraction: float
    % :param y_fraction: The fraction of the pixels used in y axis
    % :type y_fraction: float
    % :return: The image after removing background
    % :rtype: xarray DataArray
    
    x_fraction = fraction(1);
    y_fraction = fraction(2);
    offset = getBkgOffsetFromCorners(img, x_fraction, y_fraction);
    ret    = img - offset;
end

function ret = cropODImage(img, center, span)
    % Crop the image according to the region of interest (ROI).
    % :param dataSet: The images
    % :type dataSet: xarray DataArray or DataSet
    % :param center: The center of region of interest (ROI)
    % :type center: tuple
    % :param span: The span of region of interest (ROI)
    % :type span: tuple
    % :return: The cropped images
    % :rtype: xarray DataArray or DataSet

    x_start = floor(center(1) - span(1) / 2);
    x_end   = floor(center(1) + span(1) / 2);
    y_start = floor(center(2) - span(2) / 2);
    y_end   = floor(center(2) + span(2) / 2);

    ret = img(y_start:y_end, x_start:x_end);
end

function ret = calculateODImage(imageAtom, imageBackground, imageDark)
    % Calculate the OD image for absorption imaging.
    % :param imageAtom: The image with atoms
    % :type imageAtom: numpy array
    % :param imageBackground: The image without atoms
    % :type imageBackground: numpy array
    % :param imageDark: The image without light
    % :type imageDark: numpy array
    % :return: The OD images
    % :rtype: numpy array

    numerator   = imageBackground - imageDark;
    denominator = imageAtom - imageDark;

    numerator(numerator == 0)     = 1;
    denominator(denominator == 0) = 1;
    
    ret = -log(double(abs(denominator ./ numerator)));

    if numel(ret) == 1
        ret = ret(1);
    end
end

function [optrefimages] = removefringesInImage(absimages, refimages, bgmask)
    % removefringesInImage - Fringe removal and noise reduction from absorption images.
    % Creates an optimal reference image for each absorption image in a set as 
    % a linear combination of reference images, with coefficients chosen to 
    % minimize the least-squares residuals between each absorption image and 
    % the optimal reference image. The coefficients are obtained by solving a 
    % linear set of equations using matrix inverse by LU decomposition.
    %
    % Application of the algorithm is described in C. F. Ockeloen et al, Improved 
    % detection of small atom numbers through image processing, arXiv:1007.2136 (2010).
    %
    % Syntax:
    %    [optrefimages] = removefringesInImage(absimages,refimages,bgmask);
    %
    % Required inputs:
    %    absimages    - Absorption image data, 
    %                   typically 16 bit grayscale images
    %    refimages    - Raw reference image data
    %       absimages and refimages are both cell arrays containing
    %       2D array data. The number of refimages can differ from the 
    %       number of absimages.
    %
    % Optional inputs:
    %    bgmask       - Array specifying background region used, 
    %                   1=background, 0=data. Defaults to all ones.
    % Outputs:
    %    optrefimages - Cell array of optimal reference images, 
    %                   equal in size to absimages.
    %
    
    % Dependencies: none
    %
    % Authors: Shannon Whitlock, Caspar Ockeloen
    % Reference: C. F. Ockeloen, A. F. Tauschinsky, R. J. C. Spreeuw, and
    %            S. Whitlock, Improved detection of small atom numbers through 
    %            image processing, arXiv:1007.2136
    % Email: 
    % May 2009; Last revision: 11 August 2010
    
    % Process inputs
    
    % Set variables, and flatten absorption and reference images
    nimgs  = size(absimages,3);
    nimgsR = size(refimages,3);
    xdim = size(absimages(:,:,1),2);
    ydim = size(absimages(:,:,1),1);
    
    R = single(reshape(refimages,xdim*ydim,nimgsR));
    A = single(reshape(absimages,xdim*ydim,nimgs));
    optrefimages=zeros(size(absimages)); % preallocate
    
    if not(exist('bgmask','var')); bgmask=ones(ydim,xdim); end
    k = find(bgmask(:)==1);  % Index k specifying background region
    
    % Ensure there are no duplicate reference images
    % R=unique(R','rows')'; % comment this line if you run out of memory
    
    % Decompose B = R*R' using singular value or LU decomposition
    [L,U,p] = lu(R(k,:)'*R(k,:),'vector');       % LU decomposition
    
    for j=1:nimgs
        b=R(k,:)'*A(k,j);
        % Obtain coefficients c which minimise least-square residuals  
        lower.LT = true; upper.UT = true;
        c = linsolve(U,linsolve(L,b(p,:),lower),upper);
    
        % Compute optimised reference image
        optrefimages(:,:,j)=reshape(R*c,[ydim xdim]);
    end
end

function [M] = ImagingResponseFunction(B)
    x     = -128:127;
    y     = x;
    [X,Y] = meshgrid(x,y);
    R     = sqrt(X.^2+Y.^2);
    PHI   = atan2(X,Y)+pi;
    
    % Fit parameters
    tau   = B(1);
    alpha = B(2);
    S0    = B(3);
    phi   = B(4);
    beta  = B(5);
    delta = B(6);
    A     = B(7);
    C     = B(8);
    a     = B(9);
    U     = heaviside(1-R/a).*exp(-R.^2/a^2/tau^2);
    THETA = S0*(R/a).^4 + alpha*(R/a).^2.*cos(2*PHI-2*phi) + beta*(R/a).^2;
    p     = U.*exp(1i.*THETA);
    M     = A*abs((ifft2(real(exp(1i*delta).*fftshift(fft2(p)))))).^2 + C;
end

function [p] = exitPupilFunction(B)
    x     = -128:127;
    y     = x;
    [X,Y] = meshgrid(x,y);
    R     = sqrt(X.^2+Y.^2);
    PHI   = atan2(X,Y)+pi;
    
    % Fit parameters
    tau   = B(1);
    alpha = B(2);
    S0    = B(3);
    phi   = B(4);
    beta  = B(5);
    delta = B(6);
    A     = B(7);
    C     = B(8);
    a     = B(9);
    U     = heaviside(1-R/a).*exp(-R.^2/a^2/tau^2);
    THETA = S0*(R/a).^4 + alpha*(R/a).^2.*cos(2*PHI-2*phi) + beta*(R/a).^2;
    p     = U.*exp(1i.*THETA);
end

function [R, Zr] = getRadialProfile(data, radialStep)
    x = (1:size(data,2))-size(data,2)/2;
    y = (1:size(data,1))-size(data,1)/2;
    % coordinate grid:
    [X,Y] = meshgrid(x,y);
    % creating circular layers
    Z_integer = round(abs(X+1i*Y)/radialStep)+1;
    % very fast MatLab calculations:
    R  = accumarray(Z_integer(:),abs(X(:)+1i*Y(:)),[],@mean);
    Zr = accumarray(Z_integer(:),data(:),[],@mean);
end

function [RadialResponseFunc] = RadialImagingResponseFunction(C, k, kmax)
    A = heaviside(1 - k/kmax) .* exp(-C(1) * k.^4);
    W = C(2) + C(3) * k.^2 + C(4) * k.^4;
    RadialResponseFunc = 0;
    for n = -30:30
        RadialResponseFunc = RadialResponseFunc + ...
                             besselj(n, C(5) * k.^2).^2 + ...
                             besselj(n, C(5) * k.^2) .* besselj(-n, C(5) * k.^2) .* cos(2 * W);
    end
    RadialResponseFunc = C(6) * 1/2 * A .* RadialResponseFunc;
end

function R = zernike_radial(n, m, rho)
    % Compute radial part of Zernike polynomial for radial order n and azimuthal order m
    R = zeros(size(rho));
    for k = 0:(n - abs(m))/2
        coeff = (-1)^k * factorial(n - k) / ...
            (factorial(k) * factorial((n + abs(m))/2 - k) * factorial((n - abs(m))/2 - k));
        R = R + coeff * rho.^(n - 2*k);
    end
end

function Z = zernike_polynomial(n, m, X, Y)
    % Convert Cartesian coordinates to polar coordinates
    [theta, rho] = cart2pol(X, Y);
    rho(rho > 1) = 0; % Restrict to unit disk
    
    % Compute radial component
    R = zernike_radial(n, m, rho);
    
    % Compute azimuthal component (sine/cosine terms)
    if m >= 0
        Z = R .* cos(m * theta);
    else
        Z = R .* sin(-m * theta);
    end
end