function [W, cost_function] = iva_l_sos(X, varargin)
% Trung's implementation - 02/2024
%
% Implementation of IVA with the multivariate laplacian distribution.
%
% For a general description of the algorithm and its relationship with others,
% see http://mlsp.umbc.edu/jointBSS_introduction.html
%
%%% Input:
% X - data observations from K data sets, i.e. X{k}=A{k}*S{k}, where A{k}
% is an N x N unknowned invertible mixing matrix and S{k} is N x T  matrix
% with the nth row corresponding to T samples of the nth source in the kth
% dataset.  For IVA it is assumed that the source is statistically
% independent of all the sources within a dataset and exactly dependent on
% at most one source in each of the other datasets. The data, X, is a 
% 3-dimensional matrix of dimensions N x T x K. 
% 
%%% Output:
% W - the estimated demixing matrices so that ideally W{k}*A{k} = P*D{k}
% where P is any arbitrary permutation matrix and D{k} is any diagonal
% invertible (scaling) matrix.  Note P is common to all datasets; this is
% to indicate that the local permuation ambiguity between dependent sources
% across datasets should ideally be resolved by IVA.
%
% cost_function - the objective function for each iteration
%
%%% Required functions in this package:
% decouple_trick.m
% getopt.m
% logdet.m
% pca_whitening.m
%
%%% Example call:
% W = iva_l_sos(X, 'initW', W0, 'minIter', 100)
% Optional input pairs and default values are given in Params
%
%%% References
% [1]- S. Bhinge, R. Mowakeaa, V.D. Calhoun, T. Adali, "Extraction of time-varying 
%    spatio-temporal networks using parameter-tuned constrained IVA." IEEE 
%    Transactions on Medical Imaging, 2019, vol. 38, no. 7, 1715-1725 
% [2] - Mowakeaa, R., Boukouvalas, Z., Long, Q., & Adali, T. (2019). "IVA using
%    complex multivariate GGD: application to fMRI analysis."
%    Multidimensional Systems and Signal Processing, 1-20.
% [3] - E. E. Kumbasar, H. Yang, V. Trung, V. D. Calhoun, & T. Adali, 
%    "Fusion of Multitask fMRI Data with Constrained Independent Vector 
%    Analysis,"2025 59th Annual Conference on Information Sciences and Systems 
%    (CISS), Baltimore, MD, USA, 2025, pp. 1-6.
% [3] - 
%%% Copyright (C) 2024 MLSP Lab
% 
% This program is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
% 
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
% GNU General Public License for more details.
% Please make sure you include the references in any related work.
% 
% A copy of the GNU General Public License can be found at
% <https://mlsp.umbc.edu/resources.html>.
% If not, see <https://www.gnu.org/licenses/>.

[N, T, K] = size(X);

%% Gather options

% build default options structure
Params = struct( ...
    'whiten', false, ... % whitening is optional
    'initW', [], ... % initial estimates for demixing matrices in W
    'maxIter', 1000, ... % max number of iterations
    'minIter', 300, ... % min number of iterations
    'minChange', 1e-6, ... % termination threshold based on change in W
    'initStepsize', 1, ... % initial step size
    'stepsizeDecrease', 0.95, ... % decrement scale in step size if cost function increases
    'minStepsize', 1e-6, ... % minimum step size
    'checkpoints', false, ... % periodically save current results
    'updateCOV', true, ... % update the covariance estimate for every k
    'sortSCVs', true, ... % reoder SCVs from most to least ill-conditioned
    'checkpointID', '' ... % for naming the checkpoint files
);

% load in user supplied options
Params = getopt(Params, varargin{:});
beta = 0.5;
stepsize = Params.initStepsize;


% compute Kotz constant while avoiding numerical error when K is large
% and/or beta is small
% COV = p_Kotz * Sigma
log_p_Kotz = log(2) / beta + gammaln((K + 2)/2/beta) - log(K) - gammaln(K/2/beta);
p_Kotz = exp(log_p_Kotz);

if Params.whiten
    [X, V] = pca_whitening(X);
end

% Pre-compute correlation in X
Rx = cell(K, K);
for k = 1:K
    Rx{k, k} = 1 / T * (X(:, :, k) * X(:, :, k)');
    for k2 = (k + 1):K
        Rx{k, k2} = 1 / T * (X(:, :, k) * X(:, :, k2)');
        Rx{k2, k} = Rx{k, k2}';
    end
end

% Initialize W
if ~isempty(Params.initW)
    W = Params.initW;
    if Params.whiten
        for k = 1:K
            W(:, :, k) = W(:, :, k) / V(:, :, k);
        end
    end
else
    W = randn(N, N, K);
end
% make sure rows of W is normalized
for k = 1:K
    W(:, :, k) = normr(W(:, :, k));
end

% Initialize source-related quantities based on initial W
COV_N_est = zeros(K, K, N);
Y = zeros(size(X));
for k = 1:K
    Y(:, :, k) = W(:, :, k) * X(:, :, k);

    for n = 1:N
        for k2 = k:K
            COV_N_est(k, k2, n) = W(n, :, k) * Rx{k, k2} * W(n, :, k2)';
            if k < k2
                COV_N_est(k2, k, n) = COV_N_est(k, k2, n);
            end
        end
    end
end
Sigma_N_est = COV_N_est / p_Kotz;

%% Run algorithm
cost_function = [];
cost_const = N * (K / 2 * log(pi) + gammaln(K/2/beta) + K / 2 / beta * log(2) - log(beta) - gammaln(K/2));
for iter = 1:Params.maxIter

    cost = cost_const;
    for n = 1:N
        Yn = squeeze(Y(n, :, :)); % TxK
        Zn = Yn / Sigma_N_est(:, :, n); % TxK
        cost = cost + .5 * logdet(Sigma_N_est(:, :, n)) + 0.5 * mean(dot(Zn, Yn, 2).^beta);
    end
    for k = 1:K
        cost = cost - log(abs(det(W(:, :, k))));
    end
    cost_function = [cost_function, cost];
    if length(cost_function) > 1 && cost_function(end) > cost_function(end-1)
        stepsize = max(stepsize*Params.stepsizeDecrease, Params.minStepsize);
    end

    W_old = W;
    Q = 0;
    R = 0;
    for n = 1:N
        [Dn, Q, R] = decouple_trick(W, n, Q, R); % N*K

        Yn = squeeze(Y(n, :, :)); % TxK
        Zn = Yn / Sigma_N_est(:, :, n); % TxK
        ySiy = dot(Zn, Yn, 2).^(beta - 1); % Tx1
        for k = 1:K
            phi_nk = Zn(:, k) .* ySiy; % Tx1
            dWnk = (beta/T) * X(:, :, k) * phi_nk - Dn(:, k) / (W(n, :, k) * Dn(:, k)); % Nx1

            dWnk = dWnk - W(n, :, k) * dWnk * W(n, :, k)'; % Nx1
            dWnk = dWnk / norm(dWnk); % check Li's 2010 TSP paper
            temp = W(n, :, k) - stepsize * dWnk';
            W(n, :, k) = temp / norm(temp); % 1xN

            % update variables that depend on Wnk
            Yn(:, k) = W(n, :, k) * X(:, :, k);
            for k2 = 1:K
                Sigma_N_est(k, k2, n) = W(n, :, k) * Rx{k, k2} * W(n, :, k2)' / p_Kotz;
                Sigma_N_est(k2, k, n) = Sigma_N_est(k, k2, n);
            end
            if Params.updateCOV && k < K
                Zn = Yn / Sigma_N_est(:, :, n); % TxK
                ySiy = dot(Zn, Yn, 2).^(beta - 1); % Tx1
            end
        end % k

        Y(n, :, :) = Yn;
    end % n

    currentChange = 0;
    for k = 1:K
        currentChange = max(currentChange, max(1-abs(diag(W_old(:, :, k)*W(:, :, k)'))));
    end

    if currentChange < Params.minChange && iter > Params.minIter
        break;
    end

    if Params.checkpoints && mod(iter, 25) == 0
        save(['cIVA_', num2str(iter), '_', char(datetime('now'), 'MM_dd_yyyy_HH_mm_ss'), '.mat'], 'W', 'cost_function', ...
            'stepsize', 'currentChange', 'COV_N_est', 'dWnk')
    end
end

%% Dewhitening
if Params.whiten
    for k = 1:K
        W(:, :, k) = W(:, :, k) * V(:, :, k);
    end
end

% Resort order of SCVs from most to least ill-conditioned
if Params.sortSCVs
    det_COV = zeros(n, 1);
    for n = 1:N
        det_COV(n) = det(COV_N_est(:, :, n));
    end
    [~, isort] = sort(det_COV);
    % COV_N_est=COV_N_est(:,:,isort);
    for k = 1:K
        W(:, :, k) = W(isort, :, k);
    end
end

end
