function [W, cost_function, log_joint_isi] = iva_mggd(X, varargin)
% Trung's implementation - 10/2024
%
% Implementation of IVA with the multivariate generalized
% Gaussian distribution (MGGD). 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. 
% For beta=1, the distribution is Gaussian, 
% for beta<1, super-Gaussian and for beta>1, sub-Gaussian.
% In this implementation, the covariance matrix (or essentially dispersion matrix)
% can be updated less frequently by setting the parameter updateCOV to
% false. This leads to significantly faster runtime but the performance
% slightly decreases.
%
% 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 unknown 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. 
% 
% 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 NxNxK 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 permutation 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
%
%%% Difference with IVA-L-SOS implementation at http://mlsp.umbc.edu/codes/iva_l_sos_v1.m
% Cost function and gradients are computed using the dispersion matrix $\Sigma$
% instead of the covariance matrix. This ensures that the implementation is
% consistent with the algorithm's derivation. Only in the Gaussian case (beta=1), 
% the covariance matrix equals the dispersion matrix. 
% In other cases, it is important to recognize that the covariance matrix 
% is a scaled version of the dispersion matrix.
% 
% iva_mggd implements the general case for any beta>0. 
% The Kotz constant (that is the scale factor between the covariance matrix
% and the dispersion matrix) is computed in a logarithmic manner, 
% avoiding numerical issues that can happen when K goes to infinity 
% or beta goes to 0.
%
% Optimization and simplification:
% - rewrite the code for computing cost function: 
% comp_l_sos_cost function is too lengthy and hard to read. 
% An important fix is the use of logdet function which avoids numerical 
% issues when computing the logarithm of determinant of covariance matrices.
% - rewrite the code for whiten function: the new code, named "pca_whitening",
% optimizes the computation of whitening transformation by using matrix-vector
% multiplication/division instead of matrix-matrix multiplication/division.
% - rewrite the code for bss_isi function: the new code, named "joint_isi",
% simplifies the computation of Amari joint ISI by only considering the 
% most popular case: 3-D matrices. Using vectorization on the columns/rows 
% of global mixing-demixing matrix is significantly faster than the 
% previous code that uses for loop. In addition, the new implementation 
% enables partial joint ISI by allowing users to specify the number of 
% components used to compute the partial metric.
% - simplify the code for decouple_trick function: the new code only 
% includes method 4 - the QR recursive algorithm.
%
%%% Example call with multivariate Gaussian distribution (beta=1):
% W = iva_mggd(X, 'initW', W0, 'beta', 1)
% Optional input pairs and default values are given in Params
%
%%% References
% [1] - 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);

%% Gather options

% build default options structure
Params = struct( ...
    'whiten', false, ... % whitening is optional
    'beta', .5, ... % shape parameter - default multivariate Laplace distribution
    'initW', [], ... % initial estimates for demixing matrices in W
    'A', [], ... % true mixing matrices A
    '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 = Params.beta;
stepsize = Params.initStepsize;
supplyA=~isempty(Params.A); % set to true if user has supplied true mixing matrices


% 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;

%% When mixing matrix A supplied
% A matrix is conditioned by V if data is whitened
if supplyA
   if Params.whiten
      for k = 1:K
         Params.A(:,:,k) = V(:,:,k)*Params.A(:,:,k);
      end
   end
end


%% Run algorithm
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(iter) = cost;
    if length(cost_function) > 1 && cost_function(end) > cost_function(end-1)
        stepsize = max(stepsize*Params.stepsizeDecrease, Params.minStepsize);
    end

    % Some additional computations of performance via ISI when true A is supplied
    if supplyA
        log_joint_isi(iter)=joint_isi(W,Params.A);
    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);
    for k = 1:K
        W(:, :, k) = W(isort, :, k);
    end
end

end
