function [W, cost_function, eps_yr] = ar_civa_mggd(X, ref_signals, varargin)
% Trung's implementation - 10/2024
%
% Implementation of the multivariate generalized
% Gaussian distribution (MGGD) that incorporates prior information
% about the sources into the IVA cost function. 
% The shape parameter (beta) can be specified as user-input via varargin. 
% By default, iva_mggd uses beta=0.5, assuming 
% the sources follow multivariate Laplacian distribution.
%
% By alternating between a conservative scheme and an assertive scheme, 
% ar-cIVA-MGGD optimally controls the effect of each reference on 
% the corresponding estimated source. 
% There is no need for users to specify the degree of similarity 
% between the estimate and the reference signal.
%
% 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. 
% 
% ref_signals - reference signals (T x n_refs), where the number of
% references can be less than the number of sources (n_refs <= N)
%
% varargin - an optional input which is a structure containing parameters 
% of the algorithm. Default values are listed next to the structure elements.
%
%%% 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
%
% eps_yr - similarity between estimated source and corresponding reference (n_refs x K)
%
%%% Required functions in this package:
% decouple_trick.m
% getopt.m
% logdet.m
% pca_whitening.m
%
%%% Example call with multivariate Laplace distribution (beta=0.5):
% W = ar_civa_mggd(X, ref_signals, 'initW', W0, 'beta', 0.5)
% Optional input pairs and default values are given in Params
%
%%% References
% [1] - Vu, Trung, Francisco Laport, Hanlu Yang, Vince D. Calhoun, and Tulay Adali. 
% "Constrained independent vector analysis with reference for multi-subject fMRI analysis," 
% IEEE Transactions on Biomedical Engineering, 12 pages, July, 2024, 
% DOI: 10.1109/TBME.2024.3432273.
% [2] - Anderson, Matthew, Geng-Shen Fu, Ronald Phlypo, and Tülay Adalı. 
% "Independent vector analysis: Identification conditions and performance bounds." 
% IEEE Transactions on Signal Processing 62, no. 17 (2014): 4399-4410.
%
%%% 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);
[T, n_refs] = size(ref_signals);

%% Gather options

% build default options structure
Params = struct( ...
    'whiten', false, ... % whitening is optional
    'refNormalize', true, ... % normalize the references to 0 mean 1 variance
    '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.8, ... % decrement scale in step size if cost function increases
    'minStepsize', 1e-6, ... % minimum step size
    'rhoList', 0:0.01:1, ... % correlation threshold for constrained algorithm (default 0.7)
    'beta', 0.5, ... % shape parameter
    'gamma', 100, ... % penalty parameter
    'maxMu', 1, ... % maximum value of mu
    'checkpoints', false, ... % periodically save current results
    'updateCOV', true, ... % update the covariance estimate for every k
    'checkpointID', '' ... % for naming the checkpoint files
);

% load in user supplied options
Params = getopt(Params, varargin{:});
beta = Params.beta;
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

if Params.refNormalize
    ref_signals = normalize(ref_signals);
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);
dE = zeros(N, n_refs, K);
eps_yr = zeros(n_refs, K);
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
        if n <= n_refs
            dE(:, n, k) = X(:, :, k) * ref_signals(:, n) / T;
            eps_yr(n, k) = W(n, :, k) * dE(:, n, k);
        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));
alpha = randn(n_refs, K);
mu = max(0, alpha);
mu_idx = false(size(mu));
for iter = 1:Params.maxIter
    % for each n, find rho that is closest to e
    idx = arrayfun(@(e)(find(Params.rhoList > e, 1)), abs(eps_yr)); % MxK
    idx(mu_idx) = idx(mu_idx) - 1;
    rho = Params.rhoList(idx); % MxK

    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
    temp = (max(0, alpha)).^2 - mu.^2;
    constraint_cost = sum(temp(:)) / 2 / Params.gamma;
    cost_function = [cost_function, cost + constraint_cost];
    if length(cost_function) > 1 && cost_function(end-1) < cost_function(end)
        stepsize = max(stepsize*Params.stepsizeDecrease, Params.minStepsize);
    end

    W_old = W;
    Q = 0;
    R = 0;
    mu = max(0, alpha);
    mu_idx = mu_idx | (mu >= Params.maxMu);
    mu_idx = mu_idx & (mu > 0);
    mu = min(Params.maxMu, mu);
    alpha = Params.gamma * (rho - abs(eps_yr)) + mu;
    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

            if n <= n_refs
                dWnk = dWnk - max(0, alpha(n, k)) * sign(eps_yr(n, k)) * dE(:, n, k); % Nx1
            end

            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
            if n <= n_refs
                eps_yr(n, k) = W(n, :, k) * dE(:, n, k);
            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, 50) == 0
        save(['cIVA_', num2str(iter), '_', char(datetime('now'), 'MM_dd_yyyy_HH_mm_ss'), '.mat'], 'W', ...
            'cost_function', 'alpha', 'eps_yr', 'rho', 'mu', 'stepsize', 'currentChange')
    end
end

%% Clean-up Outputs
if Params.whiten
    for k = 1:K
        W(:, :, k) = W(:, :, k) * V(:, :, k);
    end
end

end
