% plot result
addpath('result')

N_per_all = 2:7;% different number of perturbations
A_name = {'n-4','n-11','n-39','n-50','n-53','n-58'};% model name
Num_non_zeros = [4,11,39,50,53,58];% number of non-zero entries in transition matrices of those 7 models
N_sparsity = length(A_name);% number of network models: 7
Size_network = 300;% number of species in the network
T1 = 5;% number of time points in each perturbation

% get results
A_error_all = cell(length(N_per_all),1);% RMSE value of transition matrix A from step 2
A_whole_cs = cell(length(N_per_all),N_sparsity);% inferred transition matrix A from step 2
A_real = cell(length(N_per_all),N_sparsity);% real transition matrix A
x_error_cell1 = cell(length(N_per_all),N_sparsity);% absolute error (inferred species in step1 - real values)
A_error_step3 = zeros(length(N_per_all),N_sparsity);% RMSE of refined A from step3
A_var_mean = zeros(length(N_per_all),N_sparsity);% variance of transition matrix from step 3

% for different number of perturbations
for i = 1:length(N_per_all)
    N_per = N_per_all(i);% the current number of perturbations
    % load results
    load(strcat('Purterbation',num2str(N_per),'Perf.mat'))
    % store RMSE of A for different models under a certain number of
    % perturnations
    A_error_all_i = zeros(N_sparsity,2);
    % for different models
    for j = 1:N_sparsity 
        result = Perf{j};
        x = result{1} ;% real x
        A = result{2};% real A
        x_pred = result{3}; % step1: predicted x
        A_pred = result{4}; % step 2: predicted A
        x_error = result{5}; % step 1: RMSE of predicted x
        A_error = result{6}; % step 2: RMSE of predicted A
        sel_ind = result{7}; % which row of A is selected in refinement step
        A_EBCS = result{8};%step3: refined A
        A_EBCS_error = result{9};%step3: RMSE of refined A
        A_EBCS_var = result{10};%step3: variance of refined A
        % get the absolute error of x in step1
        x_temp = abs(x_pred-x);% absolute error of x in Step1
        x_temp2 = x_temp(:,sum(x_temp)>0);% remove the zero error for initial state
        x_error_cell1{i,j} = x_temp2(:);% save non-zero absolute prediction error of x in Step1
        % get the RMSE of A in step2 and step3 respectively
        A_error_step3(i,j) = A_EBCS_error;% RMSE of A in step3
        A_error_all_i(j,:) = [A_error,A_EBCS_error];
     
        A_whole_cs{i,j} = A_pred;% inferred A in Step2
        A_real{i,j} = A;% real A
        
        % get the mean variance of the entries of A in step3
        A_var_i = zeros(length(sel_ind),length(sel_ind));
        for s = 1:length(A_EBCS_var)
            A_var_i(s,:) = diag(A_EBCS_var{s});% the diagonal of A_EBCS_var{s} contains the variance of sth row of A
        end
        A_var_mean(i,j) = sqrt(mean(mean(A_var_i))); % calculate the mean variance of the entries of A in step3
    end
    A_error_all{i} = A_error_all_i;
end

%% Fig. 2. Boxplots of RMSE values of inferred concentrations.
figure;
x_error_cell1_median = zeros(N_sparsity,length(N_per_all));% median value of RMSE for x
x_error_cell1_75 = zeros(N_sparsity,length(N_per_all));% the 3rd quartile value of RMSE for x

for j = 1:N_sparsity
    subplot(6,1,j)
    x_error_cell1_7 = x_error_cell1(:,j);% absolute error of x from Step1
    % find max number of species under different perturbation
    max_len = 0;
    for i  = 1:length(N_per_all)
        if length(x_error_cell1_7{i})>max_len
            max_len = length(x_error_cell1_7{i});
        end
    end
    % document median and the 3rd quartile values
    x_error_cell1_mat = zeros(max_len,length(N_per_all));
    for i  = 1:length(N_per_all)
        xi = [x_error_cell1_7{i};-1*ones(max_len-length(x_error_cell1_7{i}),1)];
        xi(xi==-1) = NaN;
        x_error_cell1_mat(:,i) = xi;
        x_error_cell1_median(j,i) = nanmedian(xi);
        x_error_cell1_75(j,i) = quantile(xi,0.75);        
    end
    % boxplots of the RMSE values of x
    boxplot(x_error_cell1_mat,'symbol','r.','outliersize',2)
    set(gca,'xTickLabel',A_name, 'FontSize',10)
    axis([0.5 6.5 -0.01 0.1])
    ylabel(A_name{j})
end
set(gca,'XTick',1:6)
per_name = {'2','3','4','5','6','7'};
set(gca,'XTickLabel',per_name, 'FontSize',10)
% text(-0.2,0.5,'Error','Rotation',90,'FontSize',10)
% text(3.1,-0.1,'Perturbations', 'FontSize',10)

% the median of the absolute error of x in step1
x_error_cell1_median %row:different network model; column: different number of perturbations
% the 3rd quartile of the absolute error of x in step1
x_error_cell1_75 %row:different network model; column: different number of perturbations

%% Fig. 3. ROC curves of network structure inference. 
A_roc = cell(length(N_per_all),N_sparsity);
A_auc = zeros(length(N_per_all),N_sparsity);
for i = 1:length(N_per_all)
    for j = 1:N_sparsity % for different models
        A_pred = A_whole_cs{i,j};% step 2: predicted A
        A = A_real{i,j};% real A
        A1 = A(:);
        A1(abs(A1)>0) = 1;% any non-zero entry represents a link
        A_pred1 = A_pred(:);
        % ROC
        [FP,TP,T,AUC] = perfcurve(A1,abs(A_pred1),1);
        A_roc{i,j} = [FP,TP];% save false positive and true positive
        A_auc(i,j) = AUC; % save AUC
    end
end
color = {'b','g','r','c','m','y','k'};
figure;
for i = 1:length(N_per_all)
    subplot(1,6,i)
    length_str = cell(1,N_sparsity);
    for j = 1:N_sparsity % for different models
        XY = A_roc{i,j};
        if j < N_sparsity
            plot(XY(:,1),XY(:,2),color{j},'LineWidth',2)
        else
            plot(XY(:,1),XY(:,2),'Color',[0.8 0.69 0.21],'LineWidth',2)
        end
        hold on
    end
    
    title(strcat(num2str(N_per_all(i)),' Perturbations'));
    axis([0 1 0.75 1])
end
legend(A_name,'Location','SouthEast')
% text(-1.2,0.7,'False positive rate')
% text(-3.5,0,'True positive rate','Rotation',90)

% the auc value of roc curves
A_auc %row:different number of perturbations; column: different network model

%% Fig. 4. Bar charts of RMSE values for inferred transition matrices. 
A_error_all2 = cell(1,N_sparsity);% store RMSE values of A inferred in step2 and step3
for i = 1:length(N_per_all)
    A_error_all_i = A_error_all{i};
    for j = 1:N_sparsity
        A_error_all2_j = A_error_all2{j};
        A_error_all2_j = [A_error_all2_j;A_error_all_i(j,:)];
        A_error_all2{j} = A_error_all2_j;
    end
end
figure;
for j = 1:N_sparsity
    subplot(6,1,j)
    bar(A_error_all2{j})
    axis([0.5 6.5 0 0.03])
    set(gca,'xTickLabel','')
    ylabel(A_name{j})
end
colormap(summer)
legend('Step2','Step3','Orientation','horizontal')
set(gca,'XTick',1:6)
per_name = {'2','3','4','5','6','7'};
set(gca,'XTickLabel',per_name, 'FontSize',10)
text(-0.2,0.11,'RMSE','Rotation',90,'FontSize',10)
text(3,-0.015,'Perturbations', 'FontSize',10)

% calculate the range of RMSE in step2 and step3
A_error_method1 = zeros(length(N_per_all),N_sparsity);
A_error_method2 = zeros(length(N_per_all),N_sparsity);
for j = 1:N_sparsity
    A_j = A_error_all2{j};
    A_error_method1(:,j) = A_j(:,1);
    A_error_method2(:,j) = A_j(:,2);
end
% the maximum, minimum and mean values of the RMSE values from step2
A_method1_range = [max(max(A_error_method1)),min(min(A_error_method1)),mean(mean(A_error_method1))];
% the maximum, minimum and mean values of the RMSE values from step3
A_method2_range = [max(max(A_error_method2)),min(min(A_error_method2)),mean(mean(A_error_method2))];
% the ratio of this two methods
A_error_method1_2 = A_error_method2./A_error_method1;
% the maximum, minimum and mean values of the ratio
A_method12_range = [max(max(A_error_method1_2)),min(min(A_error_method1_2)),mean(mean(A_error_method1_2))];

%% Fig 5. 
figure;
log_A_variance1 = log(A_var_mean(:));% log transform of the variance of A from step3
A_RMSE1 = A_error_step3(:);% RMSE of A in step3
[log_A_variance ord] = sort(log_A_variance1,'descend');
A_RMSE = A_RMSE1(ord);
p = polyfit(log_A_variance,A_RMSE,2);
f = polyval(p,log_A_variance);
plot(log_A_variance1,A_RMSE1,'o',log_A_variance,f,'b-')
zeros(length(N_per_all),N_sparsity)
for i = 1:length(N_per_all)
    for j = 1:N_sparsity
        x_ij = log(A_var_mean(i,j));
        y_ij = A_error_step3(i,j);
        text(x_ij-0.3, y_ij, strcat(num2str(N_per_all(i)),'-', A_name{j}), 'FontSize',10);
    end
end
legend('quadratic polynomial curve fitting','RMSE vs. log(variance)')
