%% 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 = "C:/Users/Karthik/Documents/GitRepositories/Calculations/Data-Analyzer/";

run        = '0013';

folderPath = strcat(folderPath, run);

cam            = 5; 

angle          = 0;
center         = [1285, 2105];
span           = [200, 200];
fraction       = [0.1, 0.1];

pixel_size     = 5.86e-6;
removeFringes  = false;

%% 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  = im2double(imrotate(h5read(fullFileName, append(groupList(cam), "/atoms")), angle));
    bkg_img  = im2double(imrotate(h5read(fullFileName, append(groupList(cam), "/background")), angle));
    dark_img = im2double(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), 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
%% Display Images

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

% Calculate the x and y limits for the cropped image
y_min = center(1) - span(2) / 2;
y_max = center(1) + span(2) / 2;
x_min = center(2) - span(1) / 2;
x_max = center(2) + span(1) / 2;

% Generate x and y arrays representing the original coordinates for each pixel
x_range = linspace(x_min, x_max, span(1));
y_range = linspace(y_min, y_max, span(2));

% Display the cropped image
for k = 1 : length(od_imgs)
    imagesc(x_range, y_range, od_imgs{k})
    axis equal tight;
    hcb = colorbar;
    hL = ylabel(hcb, 'Optical Density', 'FontSize', 16);
    set(hL,'Rotation',-90);
    colormap jet;
    set(gca,'CLim',[0 3.0]);
    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
    xlabel('Horizontal', 'Interpreter', 'tex','FontSize',16);
    ylabel('Vertical', 'Interpreter', 'tex','FontSize',16);

    drawnow
    pause(0.5)
end

%% Get rotation angles
theta_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, 'rot_mag_fin_pol_angle')
            theta_values(k) = 180 - info.Attributes(i).Value;
        end
    end
end

%% Run Fourier analysis over images

fft_imgs                    = cell(1, nimgs);

% Create VideoWriter object for movie
videoFile = VideoWriter('Single_Shot_FFT.avi', 'Motion JPEG AVI');
videoFile.FrameRate = 2;  % Set the frame rate (frames per second)
open(videoFile);  % Open the video file to write

% Display the cropped image
for k = 1 : length(od_imgs)
    IMG              = od_imgs{k};
    [IMGFFT, IMGBIN] = computeFourierTransform(IMG);
    
    figure(2);
    clf
    set(gcf,'Position',[50 50 1500 550])
    set(gca,'FontSize',16,'Box','On','Linewidth',2);
    t = tiledlayout(1, 3, 'TileSpacing', 'compact', 'Padding', 'compact'); % 1x2 grid
    
    % Calculate the x and y limits for the cropped image
    y_min = center(1) - span(2) / 2;
    y_max = center(1) + span(2) / 2;
    x_min = center(2) - span(1) / 2;
    x_max = center(2) + span(1) / 2;
    
    % Generate x and y arrays representing the original coordinates for each pixel
    x_range = linspace(x_min, x_max, span(1));
    y_range = linspace(y_min, y_max, span(2));
    
    % Display the cropped image
    nexttile
    imagesc(x_range, y_range, IMG)
    % 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
    % Top-right corner (normalized axis coordinates)
    text(1 - x_offset, 1 - y_offset, ['Angle: ', num2str(theta_values(k), '%.1f')], ...
        'Color', 'white', 'FontWeight', 'bold', 'Interpreter', 'tex', 'FontSize', 20, 'Units', 'normalized', 'HorizontalAlignment', 'right', 'VerticalAlignment', 'top');

    axis equal tight;
    hcb = colorbar;
    hL = ylabel(hcb, 'Optical Density', 'FontSize', 16);
    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
    xlabel('X', 'Interpreter', 'tex','FontSize',16);
    ylabel('Y', 'Interpreter', 'tex','FontSize',16);
    
    nexttile
    imagesc(x_range, y_range, IMGBIN)
    axis equal tight;
    hcb = colorbar;
    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
    xlabel('X', 'Interpreter', 'tex','FontSize',16);
    ylabel('Y', 'Interpreter', 'tex','FontSize',16);
    title('Denoised - Masked - Binarized','FontSize',16);
    
    nexttile
    [rows, cols] = size(IMGFFT);
    zoom_size = 50;             % Zoomed-in region around center
    mid_x = floor(cols/2);
    mid_y = floor(rows/2);
    zoomedIMGFFT = IMGFFT(mid_y-zoom_size:mid_y+zoom_size, mid_x-zoom_size:mid_x+zoom_size);
    fft_imgs{k} = zoomedIMGFFT;
    imagesc(log(1 + zoomedIMGFFT));
    % 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
    % Top-right corner (normalized axis coordinates)
    text(1 - x_offset, 1 - y_offset, ['Angle: ', num2str(theta_values(k), '%.1f')], ...
        'Color', 'white', 'FontWeight', 'bold', 'Interpreter', 'tex', 'FontSize', 20, 'Units', 'normalized', 'HorizontalAlignment', 'right', 'VerticalAlignment', 'top');
    axis equal tight;
    hcb = colorbar;
    set(gca,'YDir','normal')
    xlabel('X', 'Interpreter', 'tex','FontSize',16);
    ylabel('Y', 'Interpreter', 'tex','FontSize',16);
    title('Fourier Power Spectrum','FontSize',16);
    
    drawnow
    pause(0.5)
    
    % Capture the current frame and write it to the video
    frame = getframe(gcf);  % Capture the current figure as a frame
    writeVideo(videoFile, frame);  % Write the frame to the video
end

% Close the video file
close(videoFile);

%% Averaged FFT
% Assuming od_imgs is a cell array of size 4*n
n = length(fft_imgs) / 4;  % Calculate n
fft_imgs_avg = cell(1, n);   % Initialize the new cell array to hold the averaged images

for i = 1:n
    % Take the 4 corresponding images from od_imgs
    img1 = fft_imgs{4*i-3};  % 1st image in the group
    img2 = fft_imgs{4*i-2};  % 2nd image in the group
    img3 = fft_imgs{4*i-1};  % 3rd image in the group
    img4 = fft_imgs{4*i};    % 4th image in the group
    
    % Compute the average of these 4 images
    avg_img = (img1 + img2 + img3 + img4) / 4;
    
    % Store the averaged image in the new cell array
    fft_imgs_avg{i} = avg_img;
end

% Create VideoWriter object for movie
videoFile = VideoWriter('Averaged_FFT.avi', 'Motion JPEG AVI');
videoFile.FrameRate = 2;  % Set the frame rate (frames per second)
open(videoFile);  % Open the video file to write

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

% Display the cropped image
for k = 1 : length(fft_imgs_avg)
    imagesc(log(1 + fft_imgs_avg{k}));
    % 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
    % Top-right corner (normalized axis coordinates)
    text(1 - x_offset, 1 - y_offset, ['Angle: ', num2str(theta_values(k), '%.1f')], ...
        'Color', 'white', 'FontWeight', 'bold', 'Interpreter', 'tex', 'FontSize', 20, 'Units', 'normalized', 'HorizontalAlignment', 'right', 'VerticalAlignment', 'top');
    axis equal tight;
    hcb = colorbar;
    set(gca,'YDir','normal')
    xlabel('X', 'Interpreter', 'tex','FontSize',16);
    ylabel('Y', 'Interpreter', 'tex','FontSize',16);
    title('Averaged Fourier Power Spectrum','FontSize',16);

    drawnow
    pause(0.5)

    % Capture the current frame and write it to the video
    frame = getframe(gcf);  % Capture the current figure as a frame
    writeVideo(videoFile, frame);  % Write the frame to the video
end

% Close the video file
close(videoFile);

%% Helper Functions
function [IMGFFT, IMGBIN] = computeFourierTransform(I)
    % computeFourierSpectrum - Computes the 2D Fourier power spectrum 
    % of binarized and enhanced lattice image features, with optional central mask.
    %
    % Inputs:
    %   I           - Grayscale or RGB image matrix
    %
    % Output:
    %   F_mag       - 2D Fourier power spectrum (shifted)

    % Preprocessing: Denoise
    I_filt       = imgaussfilt(I, 1);  % adjust sigma as needed

    % Elliptical mask parameters
    [rows, cols] = size(I_filt);
    [X, Y]       = meshgrid(1:cols, 1:rows);
    cx           = cols / 2;
    cy           = rows / 2;
    
    % Shifted coordinates
    x            = X - cx;
    y            = Y - cy;
    
    % Ellipse semi-axes
    rx           = 0.4 * cols;
    ry           = 0.2 * rows;
    
    % Rotation angle in degrees -> radians
    theta_deg    = 30;        % Adjust as needed
    theta        = deg2rad(theta_deg);
    
    % Rotated ellipse equation
    cos_t        = cos(theta);
    sin_t        = sin(theta);
    
    x_rot        =  (x * cos_t + y * sin_t);
    y_rot        = (-x * sin_t + y * cos_t);
    
    ellipseMask  = (x_rot.^2) / rx^2 + (y_rot.^2) / ry^2 <= 1;
    
    % Apply cutout mask
    I_masked     = I_filt .* ellipseMask;

    % Apply global intensity threshold mask
    intensity_thresh = 0.8; 
    intensity_mask   = I_masked > intensity_thresh;
    I_masked         = I_masked .* intensity_mask;

    % Adaptive binarization
    IMGBIN           = imbinarize(I_masked, 'adaptive', 'Sensitivity', 0.0);
    
    % Compute 2D Fourier Transform
    F                = fft2(double(I));
    IMGFFT           = abs(fftshift(F))';  % Shift zero frequency to center

    % Define the radius for the circular region to exclude
    region_radius = 4;  % Adjust the radius as needed
    
    % Create a circular mask
    [~, center_idx] = max(IMGFFT(:));
    [cx, cy]                 = ind2sub(size(IMGFFT), center_idx);

    % Equation for a circle (centered at cx, cy)
    center_region            = (X - cx).^2 + (Y - cy).^2 <= region_radius^2;
    
    % Define a scaling factor for the central region (e.g., reduce amplitude by 90%)
    scaling_factor = 0.1;  % Scale center region by 10%

    % Apply the scaling factor to the center region
    IMGFFT(center_region) = IMGFFT(center_region) * scaling_factor;

end

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 ret = calculateODImage(imageAtom, imageBackground, imageDark)
    % Calculate the OD image for absorption imaging.
    % :param imageAtom: The image with atoms
    % :type imageAtom: numpy array
    % :param imageBackground: The image without atoms
    % :type imageBackground: numpy array
    % :param imageDark: The image without light
    % :type imageDark: numpy array
    % :return: The OD images
    % :rtype: numpy array

    numerator   = imageBackground - imageDark;
    denominator = imageAtom - imageDark;

    numerator(numerator == 0)     = 1;
    denominator(denominator == 0) = 1;
    
    ret = -log(double(abs(denominator ./ numerator)));

    if numel(ret) == 1
        ret = ret(1);
    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