function [X,m,S,It,Xrec,cost]=pcambtsr(X,A,maxiter,conv)
% PCAMBSR Imputes missing data in a dataset using Total Sum of Squares Reduction with PCA.
%
% SYNTAX:
% [X, m, S, It, Xrec, cost] = pcambtsr(X, A, maxiter, conv)
%
% INPUTS:
% X - Data matrix with NaNs for missing data.
% A - Number of principal components to use for imputation.
% maxiter (optional) - Maximum number of iterations for the optimization process. Default is 5000.
% conv (optional) - Convergence criterion for the optimization. Default is 1.0e-10.
%
% OUTPUTS:
% X - Original dataset with imputed values.
% m - Estimated mean vector of X.
% S - Estimated covariance matrix of X.
% It - Number of iterations performed.
% Xrec - PCA reconstruction of X with A components.
% cost - Cost function values across iterations.
%
% DESCRIPTION:
% This function imputes missing data in a dataset using Total Sum of Squares Reduction with PCA.
% It iteratively optimizes the dataset to reduce the sum of squared differences between the original and imputed data,
% using PCA to project the data into a lower-dimensional space and then reconstruct it.
% The function returns the imputed dataset, along with the estimated mean vector, covariance matrix,
% number of iterations performed, PCA reconstruction, and cost function values across iterations.
%
% Copyright (c) 2024, Odense University Hospital - Laboratory of
% Radiophysics
    
    % Check if maxiter and conv variables are provided, if not, set default values
    if ~exist('maxiter','var')
        maxiter=5000; % Default maximum number of iterations
    end
    if ~exist('conv','var')
        conv=1.0e-10; % Default convergence criterion
    end
    
    % Get the size of the data matrix X
    [n,p] = size(X);

    % Initialize a structure to hold information about missing data for each row
    pat=struct('O',[],'M',[],'nO',[],'nM',[]);
    for i=n:-1:1
        r=~isnan(X(i,:)); % Find non-NaN elements in row i
        pat(i).O = find(r==1); % Store indices of non-NaN elements
        pat(i).M = find(r==0); % Store indices of NaN elements
        pat(i).nO = size(pat(i).O,2); % Count non-NaN elements
        pat(i).nM = size(pat(i).M,2); % Count NaN elements
    end
    
    % Identify missing values and replace them with the mean of the entire dataset
    mis = isnan(X); % Find missing values
    [r, c] = find(isnan(X)); % Indices of missing values
    meanArray = ones(size(X,1),1) * mean(X,1,'omitnan'); % Array for mean values
    X(mis) = meanArray(mis); % Replace missing values with mean
    meanc=sum(X)./(n-sum(mis)); % Mean of each column, excluding missing values
    for k=1:length(r)
        X(r(k),c(k))=meanc(c(k)); % Replace missing values with column mean
    end
    
    % Initialize variables for the optimization process
    diff=100; % Initial difference for convergence check
    It=0; % Initialize iteration counter
    cost=ones(1,maxiter); % Initialize cost function values array
    
    % Optimization loop
    while It<maxiter && diff>conv
        It=It+1; % Increment iteration counter
        Xmis=X(mis); % Store original missing values
        mX=mean(X); % Calculate mean of X
        S=cov(X); % Calculate covariance matrix of X
        Xc=X-ones(n,1)*mX; % Center the data
    
        % Perform SVD on the centered data, depending on the number of rows vs. columns
        if n>p
            [~, Svals ,V]=svd(Xc,0); % SVD on the centered data
        else
            [V, Svals, ~]=svd(Xc',0); % SVD on the transpose of the centered data
        end

        singularValues = diag(Svals);
        explainedVariance = singularValues.^2 / sum(singularValues.^2) * 100; % Variance contribution per PC
        cumVariance = cumsum(explainedVariance);

        A = find(cumVariance >= 95, 1);

        V=V(:,1:A); % Select the first A principal components
    
        % Imputation step
        for i=1:n
            if pat(i).nM>0 % If there are missing values in row i
                L=V(pat(i).O,1:min(A,pat(i).nO)); % Submatrix of V for non-missing elements
                S11=S(pat(i).O,pat(i).O); % Submatrix of S for non-missing elements
                S21=S(pat(i).M,pat(i).O); % Submatrix of S for missing and non-missing elements
                z1=Xc(i,pat(i).O)'; % Transpose of non-missing elements in row i
                z2=S21*L*pinv(L'*S11*L)*L'*z1; % Calculate imputed values for missing elements
                Xc(i,pat(i).M)=z2'; % Replace missing values with calculated imputed values
            end
        end
    
        % Update X with the imputed values
        X=Xc+ones(n,1)*mX;
    
        % Calculate the difference between the original and imputed missing values
        d=(X(mis)-Xmis).^2; % Squared differences
        diff=mean(d); % Mean of squared differences
        cost(It)=diff; % Store cost function value for this iteration
    end
    
    % Post-optimization steps
    cost=cost(1:It); % Trim cost function values array
    S=cov(X); % Re-calculate covariance matrix of X
    m=mean(X); % Calculate mean of X
    [~,~,V]=svd(S,0); % Perform SVD on the covariance matrix
    P=V(:,1:A); % Select the first A principal components
    T=(X-ones(n,1)*m)*P; % Project the data onto the principal components
    Xrec=ones(n,1)*m+T*P'; % Reconstruct the data using PCA
end


