%% Extract Images
clear; close all; clc;

%% ===== D-S Settings =====
groupList          = ["/images/MOT_3D_Camera/in_situ_absorption", "/images/ODT_1_Axis_Camera/in_situ_absorption", ...
                      "/images/ODT_2_Axis_Camera/in_situ_absorption", "/images/Horizontal_Axis_Camera/in_situ_absorption", ...
                      "/images/Vertical_Axis_Camera/in_situ_absorption"];

folderPath         = "//DyLabNAS/Data/TwoDGas/2025/06/23/";

run                = '0300';

folderPath         = strcat(folderPath, run);

cam                = 5; 

angle              = 0;
center             = [1410, 2030];
span               = [200, 200];
fraction           = [0.1, 0.1];

pixel_size         = 5.86e-6;         % in meters
magnification      = 23.94;
removeFringes      = false;

ImagingMode        = 'HighIntensity';
PulseDuration      = 5e-6;            % in s

% Fourier analysis settings

% Radial Spectral Distribution
theta_min          = deg2rad(0);
theta_max          = deg2rad(180);
N_radial_bins      = 500;
Radial_Sigma       = 2;
Radial_WindowSize  = 5; % Choose an odd number for a centered moving average

% Angular Spectral Distribution
r_min              = 10;
r_max              = 20;
N_angular_bins     = 180;
Angular_Threshold  = 75;
Angular_Sigma      = 2;
Angular_WindowSize = 5;

zoom_size          = 50;             % Zoomed-in region around center

% Plotting and saving
scan_parameter              = 'ps_rot_mag_fin_pol_angle';
% scan_parameter              = 'rot_mag_field';

savefileName                = 'DropletsToStripes';
font                        = 'Bahnschrift';

if strcmp(savefileName, 'DropletsToStripes')
    scan_groups           = 0:5:45;
    titleString           = 'Droplets to Stripes';
elseif strcmp(savefileName, 'StripesToDroplets')
    scan_groups           = 45:-5:0;
    titleString           = 'Stripes to Droplets';
end

% Flags
skipNormalization           = true;
skipUnshuffling             = true;
skipPreprocessing           = true;
skipMasking                 = true;
skipIntensityThresholding   = true;
skipBinarization            = true;
skipMovieRender             = true;
skipSaveFigures             = true;
skipSaveOD                  = true;

%% ===== S-D Settings =====
groupList          = ["/images/MOT_3D_Camera/in_situ_absorption", "/images/ODT_1_Axis_Camera/in_situ_absorption", ...
                      "/images/ODT_2_Axis_Camera/in_situ_absorption", "/images/Horizontal_Axis_Camera/in_situ_absorption", ...
                      "/images/Vertical_Axis_Camera/in_situ_absorption"];

folderPath         = "//DyLabNAS/Data/TwoDGas/2025/06/24/";

run                = '0001';

folderPath         = strcat(folderPath, run);

cam                = 5; 

angle              = 0;
center             = [1410, 2030];
span               = [200, 200];
fraction           = [0.1, 0.1];

pixel_size         = 5.86e-6;         % in meters
magnification      = 23.94;
removeFringes      = false;

ImagingMode        = 'HighIntensity';
PulseDuration      = 5e-6;            % in s

% Fourier analysis settings

% Radial Spectral Distribution
theta_min          = deg2rad(0);
theta_max          = deg2rad(180);
N_radial_bins      = 500;
Radial_Sigma       = 2;
Radial_WindowSize  = 5; % Choose an odd number for a centered moving average

% Angular Spectral Distribution
r_min              = 10;
r_max              = 20;
N_angular_bins     = 180;
Angular_Threshold  = 75;
Angular_Sigma      = 2;
Angular_WindowSize = 5;

zoom_size          = 50;             % Zoomed-in region around center

% Plotting and saving
scan_parameter              = 'ps_rot_mag_fin_pol_angle';
% scan_parameter              = 'rot_mag_field';

savefileName                = 'StripesToDroplets';
font                        = 'Bahnschrift';

if strcmp(savefileName, 'DropletsToStripes')
    scan_groups           = 0:5:45
    titleString           = 'Droplets to Stripes';
elseif strcmp(savefileName, 'StripesToDroplets')
    scan_groups           = 45:-5:0;
    titleString           = 'Stripes to Droplets';
end

% Flags
skipNormalization           = true;
skipUnshuffling             = false;
skipPreprocessing           = true;
skipMasking                 = true;
skipIntensityThresholding   = true;
skipBinarization            = true;
skipMovieRender             = true;
skipSaveFigures             = true;
skipSaveOD                  = true;

%% ===== Load and compute OD image, rotate and extract ROI for analysis =====
% Get a list of all files in the folder with the desired file name pattern.

filePattern                    = fullfile(folderPath, '*.h5');
files                          = dir(filePattern);
refimages                      = zeros(span(1) + 1, span(2) + 1, length(files));
absimages                      = zeros(span(1) + 1, span(2) + 1, length(files));

for k = 1 : length(files)
    baseFileName = files(k).name;
    fullFileName = fullfile(files(k).folder, baseFileName);
    
    fprintf(1, 'Now reading %s\n', fullFileName);
    
    atm_img  = double(imrotate(h5read(fullFileName, append(groupList(cam), "/atoms")), angle));
    bkg_img  = double(imrotate(h5read(fullFileName, append(groupList(cam), "/background")), angle));
    dark_img = double(imrotate(h5read(fullFileName, append(groupList(cam), "/dark")), angle));
    
    if (isempty(atm_img)  && isa(atm_img, 'double')) || ...
       (isempty(bkg_img)  && isa(bkg_img, 'double')) || ...
       (isempty(dark_img) && isa(dark_img, 'double'))
        
        refimages(:,:,k) = nan(size(refimages(:,:,k)));  % fill with NaNs
        absimages(:,:,k) = nan(size(absimages(:,:,k)));
    else
        refimages(:,:,k)  = subtractBackgroundOffset(cropODImage(bkg_img, center, span), fraction)';
        absimages(:,:,k)  = subtractBackgroundOffset(cropODImage(calculateODImage(atm_img, bkg_img, dark_img, ImagingMode, PulseDuration), center, span), fraction)';
    end
end

% ===== Fringe removal =====

if removeFringes
    optrefimages               = removefringesInImage(absimages, refimages);
    absimages_fringe_removed   = absimages(:, :, :) - optrefimages(:, :, :);
    
    nimgs                      = size(absimages_fringe_removed,3);
    od_imgs                    = cell(1, nimgs);
    for i = 1:nimgs
        od_imgs{i}             = absimages_fringe_removed(:, :, i);
    end
else
    nimgs                      = size(absimages(:, :, :),3);
    od_imgs                    = cell(1, nimgs);
    for i = 1:nimgs
        od_imgs{i}             = absimages(:, :, i);
    end
end

% ===== Get rotation angles =====
scan_parameter_values   = zeros(1, length(files));

% Get information about the '/globals' group
for k = 1 : length(files)
    baseFileName = files(k).name;
    fullFileName = fullfile(files(k).folder, baseFileName);
    info = h5info(fullFileName, '/globals');
    for i = 1:length(info.Attributes)
        if strcmp(info.Attributes(i).Name, scan_parameter)
            if strcmp(scan_parameter, 'ps_rot_mag_fin_pol_angle')
                scan_parameter_values(k) = 180 - info.Attributes(i).Value;
            else
                scan_parameter_values(k) = info.Attributes(i).Value;
            end
        end
    end
end

% ===== Unshuffle if necessary to do so  =====

if ~skipUnshuffling
    n_values = length(scan_groups);
    n_total  = length(scan_parameter_values);
    
    % Infer number of repetitions
    n_reps   = n_total / n_values;
    
    % Preallocate ordered arrays
    ordered_scan_values = zeros(1, n_total);
    ordered_od_imgs = cell(1, n_total);
    
    counter = 1;
    
    for rep = 1:n_reps
        for val = scan_groups
            % Find the next unused match for this val
            idx = find(scan_parameter_values == val, 1, 'first');
    
            % Assign and remove from list to avoid duplicates
            ordered_scan_values(counter) = scan_parameter_values(idx);
            ordered_od_imgs{counter} = od_imgs{idx};
    
            % Mark as used by removing
            scan_parameter_values(idx) = NaN;  % NaN is safe since original values are 0:5:45
            od_imgs{idx} = [];  % empty cell so it won't be matched again
    
            counter = counter + 1;
        end
    end
    
    % Now assign back
    scan_parameter_values = ordered_scan_values;
    od_imgs = ordered_od_imgs;
end

%% Carry out PCA
numPCs = 5;

% Stack all 600 images into one data matrix [nImages x nPixels]
allImgs3D = cat(3, od_imgs{:});
[Nx, Ny]  = size(allImgs3D(:,:,1));
Xall      = reshape(allImgs3D, [], numel(od_imgs))';  % [600 x (Nx*Ny)]

% Global PCA
[coeff, score, ~, ~, explained] = pca(Xall);

%% Visualize PC1
% Extract the first principal component vector (eigenimage)
pc1_vector = coeff(:,1);  

% Reshape back to original image dimensions
pc1_image = reshape(pc1_vector, Nx, Ny);

% Plot the PC1 image
figure(1); clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
imagesc(pc1_image);
axis image off;
colormap(Colormaps.coolwarm()); % or use 'jet', 'parula', etc.
colorbar;
title(sprintf('First Principal Component (PC1) Image - Explains %.2f%% Variance', explained(1)));

%% Distribution scatter plot
numGroups = numel(scan_groups);
colors = lines(numGroups);

figure(2); clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]); hold on;
for g = 1:numGroups
    idx = scan_parameter_values == scan_groups(g);
    scatter(repmat(scan_groups(g), sum(idx),1), score(idx,1), 36, colors(g,:), 'filled');
end
xlabel('Control Parameter');
ylabel('PC1 Score');
title('Evolution of PC1 Scores');
grid on;

%% Distribution Histogram plot
numGroups = length(scan_groups);
colors = lines(numGroups);

% Define number of bins globally
numBins = 20;

% Define common bin edges based on global PC1 score range
minScore = min(score(:,1));
maxScore = max(score(:,1));
binEdges = linspace(minScore, maxScore, numBins+1);  % +1 because edges are one more than bins
binWidth = binEdges(2) - binEdges(1);  % for scaling KDE

figure(3);
clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
tiledlayout(ceil(numGroups/2), 2, 'TileSpacing', 'compact', 'Padding', 'compact');

for g = 1:numGroups
    groupVal = scan_groups(g);
    idx = scan_parameter_values == groupVal;
    groupPC1 = score(idx,1);
    
    nexttile;
    
    % Plot histogram
    histogram(groupPC1, 'Normalization', 'probability', ...
        'FaceColor', colors(g,:), 'EdgeColor', 'none', ...
        'BinEdges', binEdges);
    hold on;
    
    % Compute KDE
    [f, xi] = ksdensity(groupPC1, 'NumPoints', 1000);
    
    % Scale KDE to histogram probability scale
    f_scaled = f * binWidth;
    
    % Overlay KDE curve
    plot(xi, f_scaled, 'k', 'LineWidth', 1.5);
    
    % Vertical line at median
    med = median(groupPC1);
    yl = ylim;
    plot([med med], yl, 'k--', 'LineWidth', 1);
    
    xlabel('PC1 Score');
    ylabel('Probability');
    title(sprintf('Control Parameter = %d', groupVal));
    grid on;
    hold off;
end

sgtitle('PC1 Score Distributions');

%% Box plot for PC1 scores by group
groupLabels = cell(size(score,1),1);
for g = 1:numGroups
    idx = scan_parameter_values == scan_groups(g);
    groupLabels(idx) = {sprintf('%d', scan_groups(g))};
end

figure(4);
clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
boxplot(score(:,1), groupLabels);
xlabel('Control Parameter');
ylabel('PC1 Score');
title('Evolution of PC1 Scores');
grid on;

%% Mean and SEM plot for PC1 scores
numGroups     = length(scan_groups);
meanPC1Scores = zeros(numGroups,1);
semPC1Scores  = zeros(numGroups,1);

for g = 1:numGroups
    groupVal  = scan_groups(g);
    idx       = scan_parameter_values == groupVal;
    groupPC1  = score(idx,1);  % PC1 scores for this group
    
    meanPC1Scores(g) = mean(groupPC1);
    semPC1Scores(g)  = std(groupPC1)/sqrt(sum(idx));  % Standard error of mean
end

% Plot mean ± SEM with error bars
figure(5);
clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
errorbar(scan_groups, meanPC1Scores, semPC1Scores, 'o-', ...
    'LineWidth', 1.5, 'MarkerSize', 8, 'MarkerFaceColor', 'b');
xlabel('Control Parameter');
ylabel('Mean PC1 Score ± SEM');
title('Evolution of PC1 Scores');
grid on;

%% Plot Binder Cumulant
maxOrder = 4; % We only need up to order 4 here
numGroups = length(scan_groups);
kappa4 = NaN(1, numGroups);

for g = 1:numGroups
    groupVal = scan_groups(g);
    idx = scan_parameter_values == groupVal;
    groupPC1 = score(idx, 1);
    
    cumulants = computeCumulants(groupPC1, maxOrder);
    kappa4(g) = cumulants(4); % 4th-order cumulant
end

% Plot
figure(6);
clf; set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
plot(scan_groups, kappa4 * 1E-5, '-o', 'LineWidth', 1.5, 'MarkerFaceColor', 'b');
ylim([-12 12])
xlabel('Control Parameter');
ylabel('\kappa_4 (\times 10^{5})');
grid on;
title('Evolution of Binder Cumulant of PC1 Score');

%% --- ANOVA test ---
p = anova1(score(:,1), groupLabels, 'off');
fprintf('ANOVA p-value for PC1 score differences between groups: %.4e\n', p);

%% Helper Functions

function ret = getBkgOffsetFromCorners(img, x_fraction, y_fraction)
    % image must be a 2D numerical array
    [dim1, dim2] = size(img);

    s1 = img(1:round(dim1 * y_fraction), 1:round(dim2 * x_fraction)); 
    s2 = img(1:round(dim1 * y_fraction), round(dim2 - dim2 * x_fraction):dim2);
    s3 = img(round(dim1 - dim1 * y_fraction):dim1, 1:round(dim2 * x_fraction));
    s4 = img(round(dim1 - dim1 * y_fraction):dim1, round(dim2 - dim2 * x_fraction):dim2);

    ret = mean([mean(s1(:)), mean(s2(:)), mean(s3(:)), mean(s4(:))]);
end

function ret = subtractBackgroundOffset(img, fraction)
    % Remove the background from the image.
    % :param dataArray: The image
    % :type dataArray: xarray DataArray
    % :param x_fraction: The fraction of the pixels used in x axis
    % :type x_fraction: float
    % :param y_fraction: The fraction of the pixels used in y axis
    % :type y_fraction: float
    % :return: The image after removing background
    % :rtype: xarray DataArray
    
    x_fraction = fraction(1);
    y_fraction = fraction(2);
    offset = getBkgOffsetFromCorners(img, x_fraction, y_fraction);
    ret    = img - offset;
end

function ret = cropODImage(img, center, span)
    % Crop the image according to the region of interest (ROI).
    % :param dataSet: The images
    % :type dataSet: xarray DataArray or DataSet
    % :param center: The center of region of interest (ROI)
    % :type center: tuple
    % :param span: The span of region of interest (ROI)
    % :type span: tuple
    % :return: The cropped images
    % :rtype: xarray DataArray or DataSet

    x_start = floor(center(1) - span(1) / 2);
    x_end   = floor(center(1) + span(1) / 2);
    y_start = floor(center(2) - span(2) / 2);
    y_end   = floor(center(2) + span(2) / 2);

    ret = img(y_start:y_end, x_start:x_end);
end

function imageOD = calculateODImage(imageAtom, imageBackground, imageDark, mode, exposureTime)
%CALCULATEODIMAGE Calculates the optical density (OD) image for absorption imaging.
%
%   imageOD = calculateODImage(imageAtom, imageBackground, imageDark, mode, exposureTime)
%
%   Inputs:
%       imageAtom       - Image with atoms
%       imageBackground - Image without atoms
%       imageDark       - Image without light
%       mode            - 'LowIntensity' (default) or 'HighIntensity'
%       exposureTime    - Required only for 'HighIntensity' [in seconds]
%
%   Output:
%       imageOD         - Computed OD image
%

    arguments
        imageAtom       (:,:) {mustBeNumeric}
        imageBackground (:,:) {mustBeNumeric}
        imageDark       (:,:) {mustBeNumeric}
        mode            char {mustBeMember(mode, {'LowIntensity', 'HighIntensity'})} = 'LowIntensity'
        exposureTime    double = NaN
    end

    % Compute numerator and denominator
    numerator   = imageBackground - imageDark;
    denominator = imageAtom - imageDark;

    % Avoid division by zero
    numerator(numerator == 0)     = 1;
    denominator(denominator == 0) = 1;

    % Calculate OD based on mode
    switch mode
        case 'LowIntensity'
            imageOD = -log(abs(denominator ./ numerator));
            
        case 'HighIntensity'
            if isnan(exposureTime)
                error('Exposure time must be provided for HighIntensity mode.');
            end
            imageOD = abs(denominator ./ numerator);
            imageOD = -log(imageOD) + (numerator - denominator) ./ (7000 * (exposureTime / 5e-6));
    end

end

function [optrefimages] = removefringesInImage(absimages, refimages, bgmask)
    % removefringesInImage - Fringe removal and noise reduction from absorption images.
    % Creates an optimal reference image for each absorption image in a set as 
    % a linear combination of reference images, with coefficients chosen to 
    % minimize the least-squares residuals between each absorption image and 
    % the optimal reference image. The coefficients are obtained by solving a 
    % linear set of equations using matrix inverse by LU decomposition.
    %
    % Application of the algorithm is described in C. F. Ockeloen et al, Improved 
    % detection of small atom numbers through image processing, arXiv:1007.2136 (2010).
    %
    % Syntax:
    %    [optrefimages] = removefringesInImage(absimages,refimages,bgmask);
    %
    % Required inputs:
    %    absimages    - Absorption image data, 
    %                   typically 16 bit grayscale images
    %    refimages    - Raw reference image data
    %       absimages and refimages are both cell arrays containing
    %       2D array data. The number of refimages can differ from the 
    %       number of absimages.
    %
    % Optional inputs:
    %    bgmask       - Array specifying background region used, 
    %                   1=background, 0=data. Defaults to all ones.
    % Outputs:
    %    optrefimages - Cell array of optimal reference images, 
    %                   equal in size to absimages.
    %
    
    % Dependencies: none
    %
    % Authors: Shannon Whitlock, Caspar Ockeloen
    % Reference: C. F. Ockeloen, A. F. Tauschinsky, R. J. C. Spreeuw, and
    %            S. Whitlock, Improved detection of small atom numbers through 
    %            image processing, arXiv:1007.2136
    % Email: 
    % May 2009; Last revision: 11 August 2010
    
    % Process inputs
    
    % Set variables, and flatten absorption and reference images
    nimgs  = size(absimages,3);
    nimgsR = size(refimages,3);
    xdim = size(absimages(:,:,1),2);
    ydim = size(absimages(:,:,1),1);
    
    R = single(reshape(refimages,xdim*ydim,nimgsR));
    A = single(reshape(absimages,xdim*ydim,nimgs));
    optrefimages=zeros(size(absimages)); % preallocate
    
    if not(exist('bgmask','var')); bgmask=ones(ydim,xdim); end
    k = find(bgmask(:)==1);  % Index k specifying background region
    
    % Ensure there are no duplicate reference images
    % R=unique(R','rows')'; % comment this line if you run out of memory
    
    % Decompose B = R*R' using singular value or LU decomposition
    [L,U,p] = lu(R(k,:)'*R(k,:),'vector');       % LU decomposition
    
    for j=1:nimgs
        b=R(k,:)'*A(k,j);
        % Obtain coefficients c which minimise least-square residuals  
        lower.LT = true; upper.UT = true;
        c = linsolve(U,linsolve(L,b(p,:),lower),upper);
    
        % Compute optimised reference image
        optrefimages(:,:,j)=reshape(R*c,[ydim xdim]);
    end
end