%% Compute and plot Zernike Polynomials

NumberOfGridPoints   = 100;            % Resolution

plotZernike(2, 0, NumberOfGridPoints)  % Defocus (Z2^0)
%%
plotZernike(2, 2, NumberOfGridPoints)  % Astigmatism (Z2^2)
%%
plotZernike(3, 1, NumberOfGridPoints)  % Coma (Z3^1, x-direction)
%%
plotZernike(3, -1, NumberOfGridPoints) % Coma (Z3^-1, y-direction)
%%
plotZernike(4, 0, NumberOfGridPoints)  % Spherical Aberration (Z4^0)

%% Compute and plot Aberrated PSF, Image
NumberOfGridPoints   = 256;                         % Number of grid points per side
PupilRadius          = 10E-3;                       % Radius of pupil [m]
Length               = 32E-6;                       % Total size of the grid [m]
GridSpacing          = Length / NumberOfGridPoints; % Grid spacing [m]
Wavelength           = 421E-9;                      % Optical wavelength [m]
FocalLength          = 3*PupilRadius;               % Focal Length

% Generate PSF
C                    = [1.5, 0.0, 0.1, -0.1, 0.5];  % Zernike coefficients
PSF                  = generateAberratedPSF(C, Wavelength, PupilRadius, NumberOfGridPoints, GridSpacing, FocalLength);

xvals                = (-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * GridSpacing * 1E6;
yvals                = (-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * GridSpacing * 1E6;
plotImage(xvals, yvals, PSF, 'PSF');

% Generate object
% Object               = generateObject(Wavelength, NumberOfGridPoints, GridSpacing, FocalLength);

% Convolve the object with the PSF to simulate imaging
% Image                = convolveObjectWithPSF(abs(Object).^2, PSF);
% plotImage(xvals, yvals, Image, 'Image of Object formed by convolving with the PSF');
%% Functions

function Z = computeZernikePolynomials(n, m, r, theta)
    % Zernike polynomial function for radial and angular coordinates (r, theta)
    % Input: 
    %   n - radial order
    %   m - azimuthal frequency
    %   r - radial coordinate (normalized to unit circle, 0 <= r <= 1)
    %   theta - angular coordinate (angle in radians)
    %
    % Output:
    %   Z - Zernike polynomial value at (r, theta)
    
    if n == 2 && m == 0
        % Defocus (Z2^0)
        Z = sqrt(3) * (2 * r.^2 - 1);
    elseif n == 2 && m == 2
        % Vertical Astigmatism (Z2^2)
        Z = sqrt(6) * r.^2 .* cos(2 * theta);
    elseif n == 3 && m == 1
        % Horizontal Coma (Z3^1)
        Z = sqrt(8) * (3 * r.^3 - 2 * r) .* cos(theta);
    elseif n == 3 && m == -1
        % Vertical Coma (Z3^-1)
        Z = sqrt(8) * (3 * r.^3 - 2 * r) .* sin(theta);
    elseif n == 4 && m == 0
        % Spherical Aberration (Z4^0)
        Z = sqrt(5) * (6 * r.^4 - 6 * r.^2 + 1);
    else
        % Default to zero if no known Zernike polynomial matches
        Z = 0;
    end
end

function plotZernike(n, m, NumberOfGridPoints)
    % n: radial order
    % m: azimuthal frequency
    % NumberOfGridPoints: number of points for plotting
    
    % Create a grid of (r, theta) coordinates
    [theta, r] = meshgrid(linspace(0, 2*pi, NumberOfGridPoints), linspace(0, 1, NumberOfGridPoints));
    
    % Calculate the Zernike polynomial for the given (n, m)
    Z = computeZernikePolynomials(n, m, r, theta);
    
    % Convert polar to Cartesian coordinates for plotting
    [X, Y] = pol2cart(theta, r);
    
    % Plot the Zernike polynomial using a surface plot
    figure(1)
    clf
    set(gcf,'Position',[50 50 950 750])
    surf(X, Y, Z, 'EdgeColor', 'none');
    colormap jet;
    colorbar;
    title(sprintf('Zernike Polynomial Z_{%d}^{%d}', n, m), 'Interpreter', 'tex', 'FontSize', 16);
    xlabel('X', 'FontSize', 16);
    ylabel('Y', 'FontSize', 16);
    zlabel('Zernike Value', 'FontSize', 16);
    axis equal;
    shading interp;
end

function PSF = generateAberratedPSF(C, Wavelength, PupilRadius, NumberOfGridPoints, GridSpacing, FocalLength)
    % C is the vector of Zernike coefficients [C_defocus, C_astigmatism, C_coma, C_spherical, ...]
    % NumberOfGridPoints is the number of points for the grid (NxN grid)
    % PupilRadius is the radius of the pupil aperture
    
    % Create a grid of (x, y) of pupil-plane coordinates in Cartesian space
    [X, Y] = meshgrid((-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * GridSpacing);
    
    % Convert (x, y) to polar coordinates (r, theta)
    [theta, ~] = cart2pol(X, Y);

    % Pupil function
    P = generatePupil(Wavelength, PupilRadius, NumberOfGridPoints, GridSpacing, FocalLength);
    
    % Wavefront error from Zernike polynomials
    W = C(1) * computeZernikePolynomials(2, 0, P, theta) + ...  % Defocus (Z2^0)
        C(2) * computeZernikePolynomials(2, 2, P, theta) + ...  % Astigmatism (Z2^2)
        C(3) * computeZernikePolynomials(3, 1, P, theta) + ...  % Coma (Z3^1, x-direction)
        C(4) * computeZernikePolynomials(3, -1, P, theta) + ... % Coma (Z3^-1, y-direction)
        C(5) * computeZernikePolynomials(4, 0, P, theta);       % Spherical Aberration (Z4^0)
    
    W(P>1)       = 0;
    E            = exp(1i*2*pi*W); % Complex amplitude
    E(P>1)       = 0;              % Impose aperture size

    % Fourier transform of the pupil function with aberrations
    PSF          = abs(calculateFFT2(E)).^2;
    PSF          = PSF/sum(PSF(:));
end

function plotImage(xvals, yvals, image, titlestring)
    figure
    clf
    set(gcf,'Position',[50 50 950 750])
    imagesc(xvals, yvals, image);
    colorbar;
    colormap jet;
    title(titlestring, 'FontSize', 16);
    xlabel('X', 'FontSize', 16);
    ylabel('Y', 'FontSize', 16);
end

function P_Norm = generatePupil(Wavelength, PupilRadius, NumberOfGridPoints, GridSpacing, FocalLength)
    [X,Y]       = meshgrid((-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * ((Wavelength * FocalLength) / (NumberOfGridPoints * GridSpacing)));
    P           = sqrt(X.^2 + Y.^2);
    P_Norm      = P/PupilRadius;  % Pupil normalized by aperture radius
    assert(max(P_Norm(:))>=sqrt(2),'Sampling is not sufficient to reconstruct the entire wavefront.');
end

function G = calculateFFT2(g)
    G = fftshift(fft2(ifftshift(g)));
end

%{

function g = calculateInverseFFT2(G)
    g = ifftshift(ifft2(fftshift(G)));
end

function obj = generateObject(Wavelength, NumberOfGridPoints, GridSpacing, FocalLength)
    % Image-plane coordinates
    [U, V] = meshgrid((-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * ((Wavelength * FocalLength) / (NumberOfGridPoints * GridSpacing)));
    obj    = ;
end

function C = convolveObjectWithPSF(A, B)
    C = calculateInverseFFT2(calculateFFT2(A) .* calculateFFT2(B));
end

%}