%% STEP 1: Extract Images
clear; close all; clc;

baseDir       = 'D:/Results - Numerics/Data_Full3D/PhaseTransition/DTS/Point2';
JobNumber     = 0;
runFolder     = sprintf('Run_%03d', JobNumber);
savefileName  = 'DropletsToStripes';  % Output file name
datafileName  = './DropletsToStripes.mat';
reverseOrder  = false;  % Set this to true to reverse the theta ordering

TitleString   = 'Droplets To Stripes';

%%
baseDir       = 'D:/Results - Numerics/Data_Full3D/PhaseTransition/STD/Point2';
JobNumber     = 0;
runFolder     = sprintf('Run_%03d', JobNumber);
savefileName  = 'StripesToDroplets';  % Output file name
datafileName  = './StripesToDroplets.mat';
reverseOrder  = true;  % Set this to true to reverse the theta ordering

TitleString   = 'Stripes To Droplets';

%%
folderList       = dir(baseDir);
isValid          = [folderList.isdir] & ~ismember({folderList.name}, {'.', '..'});
folderNames      = {folderList(isValid).name};
nimgs            = numel(folderNames);

% Extract theta values from folder names
PolarAngleVals   = zeros(1, nimgs);

for k = 1:nimgs
    tokens = regexp(folderNames{k}, 'theta_(\d{3})', 'tokens');
    if isempty(tokens)
        warning('No theta found in folder name: %s', folderNames{k});
        PolarAngleVals(k) = NaN;
    else
        PolarAngleVals(k) = str2double(tokens{1}{1});
    end
end

% Choose sort direction
sortDirection    = 'ascend';
if reverseOrder
    sortDirection = 'descend';
end

% Sort folderNames based on polar angle
[~, sortIdx]     = sort(PolarAngleVals, sortDirection);
folderNames      = folderNames(sortIdx);
PolarAngleVals   = PolarAngleVals(sortIdx);  % Optional: if you still want sorted list
imgs             = cell(1, nimgs);
alphas           = zeros(1, nimgs);

for k = 1:nimgs
    folderName = folderNames{k};
    SaveDirectory = fullfile(baseDir, folderName, runFolder);

    % Extract alpha (theta) again from folder name
    tokens = regexp(folderName, 'theta_(\d{3})', 'tokens');
    alpha_val = str2double(tokens{1}{1});
    alphas(k) = alpha_val;

    matPath = fullfile(SaveDirectory, 'psi_gs.mat');
    if ~isfile(matPath)
        warning('Missing psi_gs.mat in %s', SaveDirectory);
        continue;
    end

    try
        Data = load(matPath, 'psi', 'Params', 'Transf', 'Observ');
    catch ME
        warning('Failed to load %s: %s', matPath, ME.message);
        continue;
    end

    Params = Data.Params;
    Transf = Data.Transf;
    Observ = Data.Observ;

    psi = Data.psi;
    if isgpuarray(psi)
        psi = gather(psi);
    end
    if isgpuarray(Observ.residual)
        Observ.residual = gather(Observ.residual);
    end
    
    % Axes and projection
    z = Transf.z * Params.l0 * 1e6;
    dz = z(2)-z(1);
    
    if k == 1
        x = Transf.x * Params.l0 * 1e6;
        y = Transf.y * Params.l0 * 1e6;
        % Calculate frequency increment (frequency axes)
        Nx           = length(x);       % grid size along X 
        Ny           = length(y);       % grid size along Y 
        dx           = mean(diff(x));   % real space increment in the X direction (in micrometers)
        dy           = mean(diff(y));   % real space increment in the Y direction (in micrometers)
        dvx          = 1 / (Nx * dx);  % reciprocal space increment in the X direction (in micrometers^-1)
        dvy          = 1 / (Ny * dy);  % reciprocal space increment in the Y direction (in micrometers^-1)
        
        % Create the frequency axes
        vx           = (-Nx/2:Nx/2-1) * dvx;  % Frequency axis in X (micrometers^-1)
        vy           = (-Ny/2:Ny/2-1) * dvy;  % Frequency axis in Y (micrometers^-1)
           
        % Create the Wavenumber axes
        kx_full          = 2*pi*vx;  % Wavenumber axis in X
        ky_full          = 2*pi*vy;  % Wavenumber axis in Y
    end

    n = abs(psi).^2;
    nxy = squeeze(trapz(n * dz, 3));

    imgs{k} = nxy;
end

%% STEP 2: Visualize images
figure(1);
set(gcf,'Position',[100 100 950 750])
font = 'Bahnschrift';
t = tiledlayout(4, 3, 'TileSpacing', 'compact', 'Padding', 'compact'); % 1x4 grid

for i = 1:nimgs
    ax1 = nexttile;
    imagesc(x, y, imgs{i}')
    % Define normalized positions (relative to axis limits)
    x_offset = 0.025; % 5% offset from the edges
    y_offset = 0.025; % 5% offset from the edges
    axis square;
    hcb = colorbar;
    colormap(ax1, Colormaps.plasma());
    set(gca, 'FontSize', 14);   % For tick labels only
    hL = ylabel(hcb, 'Optical Density');
    set(hL,'Rotation',-90);
    set(gca,'YDir','normal')
    % set(gca, 'YTick', linspace(y_min, y_max, 5));  % Define y ticks
    % set(gca, 'YTickLabel', flip(linspace(y_min, y_max, 5)));  % Flip only the labels
    hXLabel = xlabel('x (µm)', 'Interpreter', 'tex');
    hYLabel = ylabel('y (µm)', 'Interpreter', 'tex');
    hTitle = title('Density', 'Interpreter', 'tex');
    set([hXLabel, hYLabel, hL], 'FontName', font)
    set([hXLabel, hYLabel, hL], 'FontSize', 14)
    set(hTitle, 'FontName', font, 'FontSize', 16, 'FontWeight', 'bold');  % Set font and size for title

end

sgtitle('Ground States');

%% STEP 3: Reshape images into a 2D data matrix
% Each image will be a row vector, so size(X) = [numImages × numPixels]

% First, convert cell array to numeric 3D array
allImgs = cat(3, imgs{:});  % size: [Nx × Ny × numImages]

% Store image size for reshaping later
imgSize = size(allImgs(:,:,1));  % [Nx, Ny]

% Reshape into 2D matrix for PCA
X       = reshape(allImgs, [], nimgs)';  % [numImages × (Nx*Ny)]

%% STEP 4: Perform PCA
[coeff, score, latent, tsquared, explained] = pca(X);

% coeff     : principal component directions (eigenvectors) in pixel space
% score     : projection of each image onto the principal components
% latent    : eigenvalues (variance captured by each PC)
% explained : percentage variance explained by each PC

%% Visualize cumulative variance
figure(2);
set(gcf,'Position',[100 100 950 750])
cumVar = cumsum(explained);
plot(cumVar, 'o-', 'LineWidth', 1.5);
yline(95, '--r', '95% threshold');
xlabel('Number of Principal Components');
ylabel('Cumulative Variance (%)');
title('Minimum number of PCA components required for a basis');
grid on;

%% Visualize first few principal components
numPCsToShow = 7;
figure(3);
set(gcf,'Position',[100 100 950 750])
t = tiledlayout(3, 3, 'TileSpacing', 'compact', 'Padding', 'compact'); % 1x4 grid
for i = 1:numPCsToShow
    pcImage = reshape(coeff(:,i), imgSize); % reshape back to image form
    nexttile;
    imagesc(pcImage)
    axis image off; colormap(Colormaps.coolwarm());
    title(sprintf('PC %d (%.1f%%)', i, explained(i)));
end
sgtitle('First Principal Components (Eigenimages)');

%% Visualize evolution of PCA scores
figure(4);
set(gcf,'Position',[100 100 950 750])
plot(alphas, score(:,1), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC1 Score');
hold on;
plot(alphas, score(:,2), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC2 Score');
% plot(alphas, score(:,3), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC3 Score');
% plot(alphas, score(:,4), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC4 Score');
% plot(alphas, score(:,5), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC5 Score');
% plot(alphas, score(:,6), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC6 Score');
% plot(alphas, score(:,7), '-o', 'LineWidth', 1.5, 'DisplayName', 'PC7 Score');
legend('Location', 'northeast');
xlabel('Polar Angle (°)'); ylabel('PC Score');
title('Evolution of PC Scores');
grid on;