function [out] = NARCA(Xs, ys, Xt, yt, Ps)
%=========================================================================
%   NARCA (Version: December 22, 2023)
%   > Originally authored by Jordan Peper, Idaho State University
%   > jordanpeper@isu.edu
%
%   SUMMARY:
%   Null-augmentation regression constant analyte (NARCA) can be considered a 
%   local model updating algorithm that utilizes constant analyte target spectra to
%   update original source calibration models to better characterize the
%   deployment domain. Constant analyte spectra are obrained by mean centering
%   a collection of repeat spectra for one sample. The sample can be homgeneous or
%   heterogeneous and hence the sample can also be near-constnat analyte.
%   NARCA needs to run for each sample seperately.    
%
%   ARGUMENTS:
%   Xs: matrix containing source spectra oriented samples by wavelengths
%   ys: vector containg source reference values
%   Xt: matrix containing target spectra oriented samples by wavelengths 
%   yt: (optional) vector containing spectra oriented samples by wavelengths
%   Ps: (optional) vector containing source permutation indices 
%
%   OUTPUT:
%   > out.FULL (struct of all model statistics) & out.SEL (MDPS-selected
%     model statistics)
%       > .trgPred: predicted target labels
%       > .RMSECs: RMSE of source calibration predictions
%       > .RMSECVs: RMSE of source validation predictions
%       > .RMSEPs: RMSE of target predictions (if 'yt' provided)
%
%   > out.MISC struct of additional statistics or data
%       > .distribution: array containing selected model indices and
%         frequency
%       > .histmat: #LV by #lambdas matrix containing model selection
%         frequency at each latent variable and tuning parameter pair
%       > .LVmax: maximum latent variable used to form updated PLS models
%       > .lambdaRange: minimum and maximum scalar tunining parameter used
%         in model updating
%
%   EXAMPLE:
%   >> outNARCA = NARCA(Xs, ys, Xt, []); % Update models to target domain
%   >> ytHat = mean(mean(out.SEL.trgPred)); 
%   >> disp(ytHat); % Mean across samples, then models
%=========================================================================

%========================= IDENTIFY CONVERGENCE ==========================%
% Defines NARCA augmentation elements
Cx = Xt - mean(Xt);
Cy = zeros(size(Xt, 1), 1);

% Executes convergence method
convParams = converge(Xs, ys, Cx, Cy);
LVmax = convParams.LVmax;
lambdas = paramDecay(convParams.paramRange(1), ...
                     convParams.paramRange(2), 100);

%=========================== CROSS-VALIDATION ============================%
% Splits source data to obtain RMSECVs
[ys, isort] = sort(ys);
Xs = Xs(isort, :);
nS = round(size(Xs, 1)*0.6);
if nargin < 5 || isempty(Ps)
    Ps = round(linspace(1, nS, nS*0.6)); % 60-40 default permutation
end
Xsc = Xs(Ps, :);
ysc = ys(Ps);
Xsv = Xs; Xsv(Ps, :) = [];
ysv = ys; ysv(Ps, :) = [];

%========================== REGRESSION & STATS ===========================%
% Runs Partial Least Squares Regression (PLSR) & collects models
B = NaN(size(Xs, 2), LVmax, length(lambdas));
for j = 1:length(lambdas)
    B(:, :, j) = PLS1([Xsc - mean(Xsc); Cx*lambdas(j)], ...
                      [ysc - mean(ysc); Cy*lambdas(j)], ...
                      LVmax);
end
B = reshape(B, size(B, 1), [], 1);

% Formulates predictions
Yschats = (Xsc - mean(Xsc))*B + mean(ysc);
Ysvhats = (Xsv - mean(Xsc))*B + mean(ysc);
Ythats = (Xt - mean(Xsc))*B + mean(ysc);

% Note on dimensions: predictions matrices are m by #LVs*#lambdas; model
% statistics matrices are 1 by #LVs*#lambdas

% Calculates prediction errors & target prediciton standard deviation
RMSECs = sqrt(mean((Yschats - ysc).^2));
RMSECVs = sqrt(mean((Ysvhats - ysv).^2));
if ~isempty(yt)
    RMSEPs = sqrt(mean((Ythats - yt).^2));
else
    RMSEPs = NaN(size(RMSECs));
end
STDv = std(Ythats);

% Calculates model variances
B1norms = sum(abs(B), 1); % L1 penalty
B2norms = vecnorm(B); % L2 penalty

% Calculates target prediction differences (TPD) between each model pair
TPD = NaN(size(Ythats, 2), size(Ythats, 2));
for i = 1:size(Ythats, 2)
    for j = 1:size(Ythats, 2)
        TPD(i, j) = sum(abs(Ythats(:, i) - Ythats(:, j)));
    end
end
TPD = rangescale(TPD, 0, 1, 3);

% Calculates constant analyte U-Curve for every model pair
UCurve = rangescale((RMSECs' + RMSECs)/2)...   % Mean RMSEC
       + rangescale((RMSECVs' + RMSECVs)/2)... % Mean RMSECV
       + rangescale((B1norms' + B1norms)/2)... % Mean L1-norm
       + rangescale((B2norms' + B2norms)/2);   % Mean L2-norm

% Calculates the cosine of angle between all model pairs
CosThetas = (B' * B)./(vecnorm(B)'*vecnorm(B));

% Calculates optimization matrix
Cmat = TPD + 1*UCurve + 1.5*rangescale((STDv' + STDv)/2);

%============================ MODEL SELECTION ============================%
% Identifies the 1st percentile c-value in the diversity window
cvec = reshape(Cmat, 1, []);
costhetas = reshape(CosThetas, 1, []);
cMax = prctile(cvec(costhetas >= 0.5 & costhetas <= 0.8), 1);

% Determines indices of model pairs and their selection distribution
[indexI, indexJ] = find(CosThetas >= 0.5 & CosThetas <= 0.8 & Cmat <= cMax);
selDistribution = uniqueCounts([indexI; indexJ]');
uniqueLabels = selDistribution(:, 1);

%================================ OUTPUT =================================%
% Full statistics
out.FULL.trgPred = reshape(Ythats, LVmax, length(lambdas), []);
out.FULL.RMSECs = reshape(RMSECs, LVmax, []);
out.FULL.RMSECVs = reshape(RMSECVs, LVmax, []);
out.FULL.RMSEPs = reshape(RMSEPs, LVmax, []);

% Selected model statistics
out.SEL.trgPred = Ythats(:, uniqueLabels(:, 1));
out.SEL.RMSECs = RMSECs(uniqueLabels(:, 1));
out.SEL.RMSECVs = RMSECVs(uniqueLabels(:, 1));
out.SEL.RMSEPs = RMSEPs(uniqueLabels(:, 1));

% Miscellaneous
histmat = zeros(LVmax, length(lambdas));
histmat(uniqueLabels) = selDistribution(:, 2);
out.MISC.distribution = selDistribution;
out.MISC.histmat = histmat;
out.MISC.LVmax = LVmax;
out.MISC.lambdaRange = [lambdas(1), lambdas(2)];
end

%================================ METHODS ================================%
% Generates vector of exponentially spaced scaling tuning parameters
function params = paramDecay(mn, mx, n)
    params = exp(linspace(log(mn), log(mx), n)); % base e method
end

% Gets parameter range and max latent variable in converged tuning region
function out = converge(Xs, ys, Cx, Cy)
    % Defines augmented set
    X = [Xs - mean(Xs); Cx];

    % Initializes parameter range and max principal component
    [~, S, ~] = svd(X);
    PCmax = min([rank(X), length(find(cumsum(diag(S)/sum(diag(S))) <= 0.99))]);
    lambdas = paramDecay(1e-5, 1e5, 100);
    
    % Performs PLSR and collects RMSECs
    RMSEC = NaN(PCmax, 100);
    for k = 1:100
        B = PLS1([Xs - mean(Xs); Cx* lambdas(k)], ...
                 [ys - mean(ys); Cy * lambdas(k)], ...
                 PCmax);
        Yshats = (Xs - mean(Xs)) * B + mean(ys);
        RMSEC(:, k) = sqrt(mean((Yshats - ys).^2, 1));
    end

    % Locates converged parameters based on predetermined thresholds
    lamidx = find(abs(diff(rangescale(sum(RMSEC, 1)))) >= 0.001);
    out.paramRange = [lambdas(max(lamidx)), lambdas(min(lamidx))];
    out.LVmax = find(abs(diff(rangescale(sum(RMSEC, 2)))) >= 0.01, 1, 'last');
end

% Obtains unique labels and respective frequencies
function distribution = uniqueCounts(vec)
vec = sort(vec);
labels = unique(vec)';
counts = NaN(length(labels), 1);
for i = 1:length(labels)
    counts(i) = length(vec(vec == labels(i)));
end
distribution = [labels, counts];
end

% Range scale function
function X = rangescale(A, minBound, maxBound, direction)
    % Defualt to columnwise operation
    if nargin < 2
        minBound = 0;
        maxBound = 1;
        direction = 3;
    end
    if nargin < 4
        direction = 3;
    end
    % Row-wise operation
    if direction == 1
        for i = 1:size(A, 1)
            a = A(i, :);
            a = ((a - min(a)) * (maxBound - minBound)) / (max(a) - min(a)) + minBound;
            A(i, :) = a;
        end
    % Column-wise operation
    elseif direction == 2
        for j = 1:size(A, 2)
            a = A(:, j);
            a = ((a - min(a)) * (maxBound - minBound)) / (max(a) - min(a)) + minBound;
            A(:, j) = a;
        end
    % Matrix-wise operation
    elseif direction == 3
        A = (A - min(A, [], "all")) * (maxBound - minBound) / (max(A, [], "all") - min(A, [], "all")) + minBound;
    end
    X = A;
end

% Partial least squares (PLS) function
function [B, V] = PLS1(X, Y, k)
    X0 = X;
    Y0 = Y;
    [m, p] = size(X);
    U = zeros(m, k);
    V = zeros(p, k);
    B = zeros(p, k);
    for i = 1:k
        v = X'*Y; v = v/norm(v); V(:,i) = v;
        u = X*v;  u = u/norm(u); U(:,i) = u;
        X = X - u*(u'*X);
        Y = Y - u*(u'*Y);
        f = 1:i;
        A = U(:,f)'*X0*V(:,f);
        B(:, i) = (V(:,f)/A)*(U(:,f)'*Y0);
    end
    return
end
