% Define paths to the folders containing masked TIFF stacks
parentFolder = 'export/Upsampled_results/all_timepoints/masked_64_slices';

% Initialize a 4D array to store the pixel values
pixelValues4D = zeros(64, 682, 682, 218);

% Iterate through each folder
for i = 2:64
    % Folder name
    folderName = sprintf('Upsampled_S%d', i);
    folderPath = fullfile(parentFolder, folderName);
    
    % Initialize a 3D array to store the masked TIFF stack
    maskedStack = zeros(682, 682, 218);
    
    % Read the TIFF stack
    for j = 3:218
        filename = sprintf('masked_T%04d.tif', j);
        filePath = fullfile(folderPath, filename);
        maskedStack(:,:,j) = imread(filePath);
    end
    
    % Store the masked stack in the 4D array
    pixelValues4D(i,:,:,:) = maskedStack;
end
%% PCA

% Reshape the 4D array (64, 682, 682, 218) to a 2D array (nPixels, 218)
nSlices = size(pixelValues4D, 1);
nPixelsPerSlice = size(pixelValues4D, 2) * size(pixelValues4D, 3);
nTimePoints = size(pixelValues4D, 4);

% Reshape the 4D array to a 2D array where each row represents a pixel's time-series data
reshapedData = reshape(pixelValues4D, nSlices * nPixelsPerSlice, nTimePoints);

% Identify rows (pixels) where all values across the 218 time points are zero
nonZeroPixels = any(reshapedData ~= 0, 2);

% Exclude rows (pixels) with all zero values
filteredData = reshapedData(nonZeroPixels, :);

% Perform PCA on the filtered data
[coeff, score, ~, ~, explained] = pca(filteredData);

% 'score' contains the principal components
% We can classify the non-zero pixels based on the principal components
% %% Finding K for K-means - ELbow method
% kRange = 2:10;
% wcss = zeros(length(kRange),1);
% dataForClustering = score(:,1:10);
% 
% for kIdx = 1:length(kRange)
%     k = kRange(kIdx);
%     [~,~,sumd] = kmeans(dataForClustering,k,'Replicates',5,'Display','final');
%     
%     wcss(kIdx) = sum(sumd);
% end
% 
% figure;
% plot(kRange,wcss,'-o');
% xlabel('Number of Clusters (k)');
% ylabel('WCSS');
% title('ELbow method for optimal K');
% grid on;
% %% FInding K for K-means - Silhouette method
% kList = 1:6;  
% dataForClustering = score(:,1:10);
% 
% evaluation = evalclusters(dataForClustering, "kmeans", "silhouette", 'KList', kList);
% 
% optimalK = evaluation.OptimalK;
% fprintf('THe optimal K value based on silhouette is: %d\n',optimalK);
% 
% for kIdx = 1:length(kRange)
%     k = kRange(kIdx);
%     try
%         [clusterIdx,~] = kmeans(dataForClustering,k,'Replicates',3,'Options',opts);
%     catch ME
%         warning('K-means did not converge for k = %d. Error: %s',k,ME.message);
%         silhouetteAvg(kIdx) = NaN;
%         continue;
%     end
%     
%     silhouetteValues = silhouette(dataForClustering, clusterIdx);
%     silhouetteAvg(kIdx) = mean(silhouetteValues);
% end
% 
% figure;
% plot(kRange,silhouetteAvg,'-o');
% xlabel('Number of Clusters (k)');
% ylabel('Average Silhouette Value');
% title('Silhouette method for optimal K');
% grid on;
%% Use kmeans to cluster data points
% Use k-means clustering to classify the pixels into optimal K components
nComponents = 6;  % optimal k=4
pixelClasses = kmeans(score(:, 1:nComponents), nComponents);

% Reconstruct the classified image with zeros where necessary
classifiedImage = zeros(nSlices * nPixelsPerSlice, 1);
classifiedImage(nonZeroPixels) = pixelClasses;

% Reshape the classified image back to the original 3D spatial arrangement (64, 682, 682)
classifiedImage = reshape(classifiedImage, nSlices, size(pixelValues4D, 2), size(pixelValues4D, 3));

% % Select a slice to show, and use squeeze to make sure it's a 2D matrix
% sliceToShow = 20; % Example slice index
% classifiedSlice = squeeze(classifiedImage(sliceToShow, :, :));
% 
% % Display the classified slice as a grayscale image
% imshow(classifiedSlice, []);
% title(sprintf('Classified Slice %d', sliceToShow));
%% Explain each component
% Perform PCA on the filtered data
[coeff, score, ~, ~, explained] = pca(filteredData);

% Display the amount of variance explained by each component
figure;
plot(cumsum(explained), '-o');
xlabel('Number of Components');
ylabel('Cumulative Variance Explained (%)');
title('Variance Explained by Principal Components');

% Display the top components to see their individual contribution
disp('Variance explained by each principal component:');
disp(explained);

% Determine how many components to retain
nComponentsToKeep = find(cumsum(explained) > 98, 1); % Keep components that explain 95% variance
fprintf('Number of components to retain (98%% variance explained): %d\n', nComponentsToKeep);
        
% Plot the first few principal components to observe their trends
nComponentsToPlot = min(nComponentsToKeep, 5);  % Plot only the first few components

figure;
for i = 1:nComponentsToPlot
    subplot(nComponentsToPlot, 1, i);
    plot(coeff(:, i));  % Plot the i-th principal component
    title(sprintf('Principal Component %d', i));
    xlabel('Time Points');
    ylabel('Component Value');
end
%% % Create a folder to save the classified slices
outputFolder = 'PCA-6groups-final';
if ~exist(outputFolder, 'dir')
    mkdir(outputFolder);
end

% Loop through each slice and save it as a TIFF file
for sliceIdx = 2:64
    % Extract the classified slice
    classifiedSlice = squeeze(classifiedImage(sliceIdx, :, :));

    % Define the output file name
    outputFileName = fullfile(outputFolder, sprintf('Classified_S%03d.tif', sliceIdx));
    
    % Save the classified slice as a TIFF file
    imwrite(uint8(classifiedSlice), outputFileName);
end

disp('All classified slices have been saved as TIFF files.');

