%% This script generates the simulations presented in Figure 4 of the
% manuscript

% Adds the simulation functions folder
addpath('./LogisticRegressionFramework/Simulations');

% Number of time points
T = 1200;

% Number of hub regions
N_hubs = 0;

% Ratio between cross-validation and training dataset sizes
subject_ratio = 3/5;

% Number of networks, and of regions in each
N = 7;
N_regions_pernetwork = [5,4,7,6,4,5,4];

% Number of regions overall
R = sum(N_regions_pernetwork + N_hubs);

% Ground truth intrinsic transition probabilities
Alpha_GT = repmat(0.5,N,1);

% Declaration of co-activation stuff
Gamma_GT = zeros(R,R);
Network_indices = zeros(R,R);

% Building co-activation coefficients
idx_r = 1;

Network_assignment = nan(R,2);

for i = 1:N
    Gamma_GT(idx_r:idx_r+N_regions_pernetwork(i)-1,idx_r:idx_r+N_regions_pernetwork(i)-1) = 1;
    Network_indices(idx_r:idx_r+N_regions_pernetwork(i)-1,idx_r:idx_r+N_regions_pernetwork(i)-1) = i;
    Network_assignment(idx_r:idx_r+N_regions_pernetwork(i)-1,1) = i;
    idx_r = idx_r + N_regions_pernetwork(i);
end

for i = 1:size(Gamma_GT,1) 
    Gamma_GT(i,i) = 0;
end

tmp = diag(Network_indices);

for n = 1:N_hubs
    N1 = Hub_networks(n,1);
    N2 = Hub_networks(n,2);

    idx1 = find(tmp==N1);
    idx2 = find(tmp==N2);
    
    Gamma_GT(idx_r,idx1) = 1;
    Network_indices(idx_r,idx1) = N1;
    
    Gamma_GT(idx1,idx_r) = 1;
    Network_indices(idx1,idx_r) = N1;
    
    Gamma_GT(idx_r,idx2) = 1;
    Network_indices(idx_r,idx2) = N2;
    
    Gamma_GT(idx2,idx_r) = 1;
    Network_indices(idx2,idx_r) = N2;
    
    Network_assignment(idx_r,1) = N1;
    Network_assignment(idx_r,2) = N2;
    
    idx_r = idx_r + 1;
end

for i = 1:size(Gamma_GT,1) 
    Gamma_GT(i,i) = 0;
end

% Computes the density of coefficients for each region and on the whole
Density_coact = sum(Gamma_GT)/(R-1);
Overall_density_coact = sum(sum(Gamma_GT))/(R*(R-1));

% Declaration of the causal ground truth at the network level
Beta_GT = [0 0 0 0 0 0 0; 0 0 0 0 0 -1 0; 0 0 0 0 0 1 0; 0 0 1 0 0 0 0; ...
    0 0 0 -1 0 0 0; 0 0 0 0 0 0 0; 0 0 0 0 0 1 0];

% Now I want to enlarge this to the full regional view
Beta_GT_regions = zeros(R,R);

for r1 = 1:R
    for r2 = 1:R
        tmp1 = Network_assignment(r1,:);
        tmp2 = Network_assignment(r2,:);

        if sum(isnan(tmp1)) == 1
            if sum(isnan(tmp2)) == 1
                Beta_GT_regions(r1,r2) = Beta_GT(tmp1(1),tmp2(1));
            end
        end
    end
end

% Associated coefficient densities
Density_causal = sum(abs(Beta_GT_regions))/(R-1);
Overall_density_causal = sum(sum(abs(Beta_GT_regions)))/(R*(R-1));

% Number of training subjects at most (in the paper, we used 80)
n_training_init = 80;

% Associated number of cross-validation subjects
n_CV_init = round(subject_ratio*n_training_init);

% The data is created for each level of noise investigated; the strategy is
% to lower the number of simulated subjects within an inner loop to span
% both the noise and dataset size ranges of interest
for Noise = [sqrt(1),sqrt(2),sqrt(4),sqrt(9),sqrt(16),sqrt(25)]
    
    TC = RSCHMM_Simulate_Newlook(Alpha_GT,Alpha_GT,Beta_GT,Beta_GT,...
        Network_assignment,Noise,n_training_init,T,[],[]);
    CV = RSCHMM_Simulate_Newlook(Alpha_GT,Alpha_GT,Beta_GT,Beta_GT,...
        Network_assignment,Noise,n_CV_init,T,[],[]);
    
    save(['SIMULATION_FIGURE4_SIG2=',num2str(Noise^2),'_N=',num2str(n_training_init)],'TC','CV');

    for n_training = [50,40,30,15,1]
        
        n_CV = round(subject_ratio*n_training);
        
        TC = TC(1:n_training);
        CV = CV(1:n_CV);
        
        save(['SIMULATION_FIGURE4_SIG2=',num2str(Noise^2),'_N=',num2str(n_training)],'TC','CV');
    end
end