%% Seidel Aberration OPD Visualization
clear; close all; clc;

%% Parameters
% Aberration coefficients (waves)
W040 = 1;    % Spherical
W131 = 1;    % Coma
W222 = 1;    % Astigmatism
W220 = 1;    % Field Curvature
W311 = 1;    % Distortion

x0 = 1;      % Normalized off-axis field (0 for on-axis)
pupil_diameter = 1; % mm
radius_mm = pupil_diameter / 2;
N = 201;       % grid resolution
lambda = 0.532; % microns (optional, can scale OPD)

%% Grid in mm
x_mm = linspace(-radius_mm, radius_mm, N);
y_mm = linspace(-radius_mm, radius_mm, N);
[X, Y] = meshgrid(x_mm, y_mm);
rho_norm = sqrt(X.^2 + Y.^2)/radius_mm;
theta = atan2(Y, X);
mask = rho_norm <= 1;

%% Compute OPDs
OPD_spherical = W040 * rho_norm.^4;
OPD_spherical(~mask) = NaN;

OPD_coma = W131 * x0 * rho_norm.^3 .* cos(theta);
OPD_coma(~mask) = NaN;

OPD_astig = W222 * x0^2 * rho_norm.^2 .* cos(theta).^2;
OPD_astig(~mask) = NaN;

OPD_fieldcurv = W220 * x0^2 * rho_norm.^2;
OPD_fieldcurv(~mask) = NaN;

OPD_distortion = W311 * x0^3 * rho_norm .* cos(theta);
OPD_distortion(~mask) = NaN;

%% Store in a struct for iteration
Aberrations = struct('name', {'Spherical','Coma','Astigmatism','Field Curvature','Distortion'},...
                     'OPD', {OPD_spherical, OPD_coma, OPD_astig, OPD_fieldcurv, OPD_distortion});

%% Loop to plot each aberration
for k = 1:length(Aberrations)
    OPD = Aberrations(k).OPD;
    
    % 1) Tangential OPD Fan (x-axis, theta=0)
    rho_line = linspace(-1, 1, N);
    switch Aberrations(k).name
        case 'Spherical'
            OPD_line = W040 * rho_line.^4;
        case 'Coma'
            OPD_line = W131 * x0 * rho_line.^3;
        case 'Astigmatism'
            OPD_line = W222 * x0^2 * rho_line.^2;
        case 'Field Curvature'
            OPD_line = W220 * x0^2 * rho_line.^2;
        case 'Distortion'
            OPD_line = W311 * x0^3 * rho_line;
    end
    
    figure;
    plot(rho_line, OPD_line, 'LineWidth', 2);
    xlabel('Normalized Pupil Coordinate (\rho)');
    ylabel('OPD (waves)');
    title(['Tangential OPD Fan - ', Aberrations(k).name]);
    grid on;
    
    % 2) 2D OPD Map
    figure;
    imagesc(x_mm, y_mm, OPD);
    axis equal tight;
    colormap jet; colorbar;
    xlabel('X (mm)'); ylabel('Y (mm)');
    title(['2D OPD Map - ', Aberrations(k).name]);
    
    % 3) 3D Wavefront Surface
    figure;
    surf(X, Y, OPD, 'EdgeColor', 'none');
    axis equal;
    colormap jet; colorbar;
    xlabel('X (mm)'); ylabel('Y (mm)'); zlabel('OPD (waves)');
    title(['3D Wavefront - ', Aberrations(k).name]);
    view(30, 30);
end
