%% Parameters

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/07/16/";

run            = '0002';

folderPath     = strcat(folderPath, run);

cam            = 5; 

angle          = 0;
center         = [1430, 2025];
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;

% Plotting and saving
scan_parameter              = 'evap_rot_mag_field';
scan_groups                 = 0:10:50;
savefileName                = 'Droplets';
font                        = 'Bahnschrift';

% Flags
skipUnshuffling             = 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));
    
    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

% ===== 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, '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

%% Display Images

figure(1)
clf
set(gcf,'Position',[50 50 950 750])

% Get image size in pixels
[Ny, Nx] = size(od_imgs{1});

% Define pixel size and magnification (if not already defined earlier)
dx = pixel_size / magnification;  % e.g. in meters
dy = dx;  % assuming square pixels

% Define x and y axes in μm (centered at image center)
x = ((1:Nx) - ceil(Nx/2)) * dx * 1E6;  % micrometers
y = ((1:Ny) - ceil(Ny/2)) * dy * 1E6;

% Display the cropped image
for k = 1 : length(od_imgs)
    imagesc(x, y, od_imgs{k});
    hold on;

    % Convert pixel grid to µm (already done: x and y axes)
    % Draw ↘ diagonal (top-left to bottom-right)
    drawODOverlays(x(1), y(1), x(end), y(end));
    
    % Draw ↙ diagonal (top-right to bottom-left)
    drawODOverlays(x(end), y(1), x(1), y(end));
    
    hold off;
    axis equal tight;
    colormap(Colormaps.inferno());
    set(gca, 'FontSize', 14, 'YDir', 'normal');
    
    if strcmp(scan_parameter, 'rot_mag_fin_pol_angle')
        text(0.975, 0.975, [num2str(scan_parameter_values(k), '%.1f^\\circ')], ...
            'Color', 'white', 'FontWeight', 'bold', 'FontSize', 24, ...
            'Interpreter', 'tex', 'Units', 'normalized', ...
            'HorizontalAlignment', 'right', 'VerticalAlignment', 'top');
    else
        text(0.975, 0.975, [num2str(scan_parameter_values(k), '%.2f'), ' G'], ...
            'Color', 'white', 'FontWeight', 'bold', 'FontSize', 24, ...
            'Interpreter', 'tex', 'Units', 'normalized', ...
            'HorizontalAlignment', 'right', 'VerticalAlignment', 'top');
    end
    
    colorbarHandle = colorbar;
    ylabel(colorbarHandle, 'Optical Density', 'Rotation', -90, 'FontSize', 14, 'FontName', font);

    xlabel('x (\mum)', 'Interpreter', 'tex', 'FontSize', 14, 'FontName', font);
    ylabel('y (\mum)', 'Interpreter', 'tex', 'FontSize', 14, 'FontName', font);
    title('OD Image', 'FontSize', 16, 'FontWeight', 'bold', 'Interpreter', 'tex', 'FontName', font);
    
    drawnow;
    pause(0.5);
end


%% 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 drawODOverlays(x1, y1, x2, y2)
    
    % Parameters
    tick_spacing = 10;   % µm between ticks
    tick_length = 2;     % µm tick mark length
    line_color = [0.5 0.5 0.5];
    tick_color = [0.5 0.5 0.5];
    font_size = 10;

    % Vector from start to end
    dx = x2 - x1;
    dy = y2 - y1;
    L = sqrt(dx^2 + dy^2);

    % Unit direction vector along diagonal
    ux = dx / L;
    uy = dy / L;

    % Perpendicular unit vector for ticks
    perp_ux = -uy;
    perp_uy = ux;

    % Midpoint (center)
    xc = (x1 + x2) / 2;
    yc = (y1 + y2) / 2;

    % Number of positive and negative ticks
    n_ticks = floor(L / (2 * tick_spacing));

    % Draw main diagonal line
    plot([x1 x2], [y1 y2], '--', 'Color', line_color, 'LineWidth', 1.2);

    for i = -n_ticks:n_ticks
        d = i * tick_spacing;
        xt = xc + d * ux;
        yt = yc + d * uy;

        % Tick line endpoints
        xt1 = xt - 0.5 * tick_length * perp_ux;
        yt1 = yt - 0.5 * tick_length * perp_uy;
        xt2 = xt + 0.5 * tick_length * perp_ux;
        yt2 = yt + 0.5 * tick_length * perp_uy;

        % Draw tick
        plot([xt1 xt2], [yt1 yt2], '--', 'Color', tick_color, 'LineWidth', 1);

        % Label: centered at tick, offset slightly along diagonal
        if d ~= 0
            text(xt, yt, sprintf('%+d', d), ...
                'Color', tick_color, ...
                'FontSize', font_size, ...
                'HorizontalAlignment', 'center', ...
                'VerticalAlignment', 'bottom', ...
                'Rotation', atan2d(dy, dx));
        end

    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
