%% Zernike Polynomials

resolution   = 100;            % Resolution

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

%% Aberrated PSF
C            = [0.0, 0.0, 0.0, 0.0, 0.0]; % Zernike coefficients
pupil_radius = 0.5;                            % Pupil radius (normalized)

plotAberratedPSF(C, pupil_radius, resolution);

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, resolution)
    % n: radial order
    % m: azimuthal frequency
    % resolution: number of points for plotting
    
    % Create a grid of (r, theta) coordinates
    [theta, r] = meshgrid(linspace(0, 2*pi, resolution), linspace(0, 1, resolution));
    
    % 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 [theta, r, PSF] = modelPSF(C, pupil_radius, resolution)
    % C is the vector of Zernike coefficients [C_defocus, C_astigmatism, C_coma, C_spherical, ...]
    % resolution is the number of points for the grid (NxN grid)
    % pupil_radius is the radius of the pupil aperture
    
    % Create a grid of (r, theta) coordinates
    [theta, r] = meshgrid(linspace(0, 2*pi, resolution), linspace(0, 1, resolution));
    
    % Pupil function: 1 inside the pupil radius, 0 outside
    P = double(r <= pupil_radius);  % 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(fftshift(fft2(P .* exp(1i * W)))).^2;  % Intensity distribution
end

function plotAberratedPSF(C, pupil_radius, resolution)
    % C: Zernike coefficients [C_defocus, C_astigmatism, C_coma_x, C_coma_y, C_spherical]
    % pupil_radius: Radius of the pupil aperture
    % resolution: Number of points for plotting
    
    % Generate PSF using the updated modelPSF function
    [theta, r, PSF] = modelPSF(C, pupil_radius, resolution);
    
    % Convert polar to Cartesian coordinates for plotting
    [X, Y] = pol2cart(theta, r);
    
    figure(1)
    clf
    set(gcf,'Position',[50 50 950 750])
    surf(X, Y, PSF, 'EdgeColor', 'none');
    view(2);  % 2D view
    axis equal tight;
    shading interp;
    colorbar;
    colormap jet;
    title('PSF', 'FontSize', 16);
    xlabel('X', 'FontSize', 16);
    ylabel('Y', 'FontSize', 16);
end