%% 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   = 1024;                        % Number of grid points per side
PupilRadius          = 0.010;                       % Radius of pupil [m]
Length               = 0.5;                         % Total size of the grid [m]
GridSpacing          = Length / NumberOfGridPoints; % Grid spacing [m]
Wavelength           = 421e-9;                      % Optical wavelength [m]
ImageDistance        = 0.7;                         % Image distance [m]

% Generate PSF
C                    = [0.3, 0.2, 0.1, -0.1, 0.4];  % Zernike coefficients
[X, Y, PSF]          = generateAberratedPSF(C, PupilRadius, NumberOfGridPoints, GridSpacing);
plotAberratedPSF(X, Y, PSF);

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

% Convolve the object with the PSF to simulate imaging
Image                = convolveObjectWithPSF(abs(Object).^2, PSF, 1);

%% 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 = 2 * r.^2 - 1;
    elseif n == 2 && m == 2
        % Astigmatism (Z2^2)
        Z = r.^2 .* cos(2 * theta);
    elseif n == 3 && m == 1
        % Coma (Z3^1)
        Z = (3 * r.^3 - 2 * r) .* cos(theta);
    elseif n == 3 && m == -1
        % Coma (Z3^-1)
        Z = (3 * r.^3 - 2 * r) .* sin(theta);
    elseif n == 4 && m == 0
        % Spherical Aberration (Z4^0)
        Z = 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 [X, Y, PSF] = generateAberratedPSF(C, PupilRadius, NumberOfGridPoints, GridSpacing)
    % 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, r] = cart2pol(X, Y);

    % Pupil function: 1 inside the pupil radius, 0 outside
    P = generateCircularMask(X, Y, 2 * PupilRadius);  % 2D mask
    
    % Wavefront error from Zernike polynomials
    W = C(1) * computeZernikePolynomials(2, 0, r, theta) + ...  % Defocus (Z2^0)
        C(2) * computeZernikePolynomials(2, 2, r, theta) + ...  % Astigmatism (Z2^2)
        C(3) * computeZernikePolynomials(3, 1, r, theta) + ...  % Coma (Z3^1, x-direction)
        C(4) * computeZernikePolynomials(3, -1, r, theta) + ... % Coma (Z3^-1, y-direction)
        C(5) * computeZernikePolynomials(4, 0, r, theta);       % Spherical Aberration (Z4^0)
    
    % Fourier transform of the pupil function with aberrations
    PSF = abs(calculateExtendedFFT2(P .* exp(-1i * 2*pi * W), GridSpacing)).^2;
end

function plotAberratedPSF(X, Y, PSF)
    % C: Zernike coefficients [C_defocus, C_astigmatism, C_coma_x, C_coma_y, C_spherical]
    % PupilRadius: Radius of the pupil aperture
    % NumberOfGridPoints: Number of points for plotting
    % Generate PSF using the updated modelPSF function
    PSF         = PSF / max(PSF(:));
    figure(1)
    clf
    set(gcf,'Position',[50 50 950 750])
    surf(X, Y, PSF, 'EdgeColor', 'none');
    xlim([-0.05 0.05])
    ylim([-0.05 0.05])
    view(2);  % 2D view
    shading interp;
    colorbar;
    colormap jet;
    title('PSF', 'FontSize', 16);
    xlabel('X', 'FontSize', 16);
    ylabel('Y', 'FontSize', 16);
end

function C = convolveObjectWithPSF(A, B, GridSpacing)
    N = size(A, 1);
    C = calculateExtendedInverseFFT2(calculateExtendedFFT2(A, GridSpacing) .* calculateExtendedFFT2(B, GridSpacing), 1/(N*GridSpacing));
end

function z = generateCircularMask(x, y, D)
    r         = sqrt(x.^2+y.^2);
    z         = double(r<D/2);
    z(r==D/2) = 0.5;
end

function G = calculateExtendedFFT2(g, Delta)
    G = fftshift(fft2(fftshift(g))) * Delta^2;
end

function g = calculateExtendedInverseFFT2(G, Delta_f)
    N = size(G, 1);
    g = ifftshift(ifft2(ifftshift(G))) * (N * Delta_f)^2;
end

function ret = generateRectangle(x, D)
    if nargin == 1, D = 1; end
     x = abs(x);
     ret = double(x<D/2);
    ret(x == D/2) = 0.5;
end

function obj = generateObject(Wavelength, ImageDistance, NumberOfGridPoints, GridSpacing)
    % Image-plane coordinates
    [U, V] = meshgrid((-NumberOfGridPoints/2 : NumberOfGridPoints/2-1) * ((Wavelength * ImageDistance) / (NumberOfGridPoints * GridSpacing)));
    obj    = (generateRectangle((U-1.4e-4)/5e-5) + generateRectangle(U/5e-5)+ generateRectangle((U+1.4e-4)/5e-5)) .* generateRectangle(V/2e-4);
end