%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Implementation of complete model order selection using IVA (CMI-IVA)
% Developed by  M.A.B.S. Akhonda
% Coded by M.A.B.S. Akhonda (mo32 at umbc.edu)
% 
% Reference:
%  M. A. B. S. Akhonda, B. Gabrielson, V. D. Calhoun and T. Adali,
%  "Complete Model Identification Using Independent Vector Analysis: Application to the Fusion of Task FMRI Data," 2021 
%  IEEE Data Science and Learning Workshop (DSLW), Toronto, ON, Canada, 2021, pp. 1-6, doi: 10.1109/DSLW51110.2021.9523414.


function [Order,IVA_results,est_results] = order_selection_IVA(X,num_comp)
%Input:
% X = Cell containig K datasets (M_k*T), where M_k  = number of obervations(Subjects or timepoints), 
%     T = Samples (voxels)
% num_comp = Initial PCA order (can be a large number or M/2)

%Output:
% Order: Cell containing (1*2) common subspace order estimated using information theoratic
% criterion (MDL and AIC)
% IVA_results: Struct containing common and distinct components and subject covariations estimated by IVA. 
% Est_results: Struct containing souce component vectors and the eigen
% properties of SCVs



K = length(X); % Number of modalities/datasets
[T] = size(X{1},2); % Number of samples 
N = num_comp; % Number of components for dimensionality reduction

%% Dimensionality reduction using SVD
    for i = 1:K
        [U2{i}, D2{i}, V2{i}]=svd(X{i},'econ');
        u_temp = U2{i};d_temp = D2{i};v_temp = V2{i};
        x_d(:,:,i)=(inv(d_temp(1:num_comp,1:num_comp)))*u_temp(:,1:num_comp)'*X{i};
    end
%end
clear U_temp D_temp V_temp

%% IVA step to maximize the correlation across modalities 
w_iva = iva_second_order_cmi_iva(x_d); % second order statistics used, replace it to use HOS

%% Subject variations and SCV estimation
% Subject covariations estimation through back reconstruction  
for i = 1:K
    u_temp = U2{i};d_temp = D2{i};v_temp = V2{i};
    a_hat{i} = (pinv(u_temp(:,1:num_comp)')*(d_temp(1:num_comp,1:num_comp)))*pinv(w_iva(:,:,i));
end
%Components and SCV estimation 
for i = 1:K
    s_iva(:,:,i) = w_iva(:,:,i)*x_d(:,:,i);
end

for n = 1:N
    for k = 1:K
        SCV(k,:,n) = s_iva(n,:,k);
    end
end

for i = 1:N
    temp = SCV(:,:,i)';
    cov_a(:,:,i) = cov(SCV(:,:,i)');
end

test = zeros(N*K,N*K); % creating N*K,N*K covariance matrix
for i = 1:N
    c1 = ((i-1)*K)+1:K*i;
    test(c1,c1) = cov_a(:,:,i);
end

temp = (test); % covariance matrix thresolding. Default value 0.3
for kk = 1:(N*K)
    Z = find(abs(test(:,kk))<0.3);
    temp(Z,kk) = 0;
end

%% Second level PCA to identify the common order
[U,D,V] = svd(temp,'econ');
est_results.U = U;
est_results.D = diag(D);
est_results.cov = cov_a;

%% ITC on estimated eigenvalues to identify the dominant ones
lamda = (diag(D)).^2;
for k = 0:N-1
    coef=1/(N-k);
    den=coef*sum(lamda(k+1:N));
    num=prod(lamda(k+1:N)).^(coef);
    penalty(k+1) =(0.5*log(T)*(k*(2*N-k)+1));
    likelihood(k+1) = T*log(((num/den)^(N-k))); %%likelihood
    aic(k+1)=-likelihood(k+1)+(k*(2*N-k)+1); %%min calculation using AIC
    mdl(k+1)=-likelihood(k+1)+(0.5*log(T)*(k*(2*N-k)+1));%%min calculation using MDL
end
%% Common order results
[value,ITC1] = min(mdl);
ITC.mdl = ITC1-1;
[value,ITC2] = min(aic);
ITC.AIC = ITC2-1;
[ITC3, idx_of_result] = knee_pt(diag(D));
ITC.knee = ITC3-1;

%% Common structure identification using eigenvectors
if ITC.mdl == 0
    count = find(round(lamda)>1);
    ITC.mdl = length(count);
end

for e = 1: ITC.mdl
    Non_zero_ele{e} = find(U(:,e));
    SCV_id(e) = ceil(mean(Non_zero_ele{e}./K));
end
store_results = zeros(ITC.mdl,K);
for e_1 = 1: ITC.mdl
    temp  = Non_zero_ele{e_1};
    dev = floor(mean(temp./K))*K;
    for t = 1:K
       check_temp = dev+t;
       if any(temp(:)==check_temp)==1
           store_results(e_1,t) = 1;
       else
           store_results(e_1,t) = 0;
       end
    end
end
%% results
Order = ITC;
IVA_results.W = w_iva;
IVA_results.S = s_iva;
IVA_results.A = a_hat;
est_results.SCV_id = SCV_id;
est_results.store = store_results;
end