clear all;
addpath(genpath('helpers'));

runs = 5;

%% set hypergraph parameters
n = 400; % number of vertices (of individuals)
K = 5; % maximum size of a hyperedge

%% set other parameters
beta = [0:0.1:0.2];
%beta = 0.2-beta;% rates of infection for the hypergraph model
delta = 1; % recovery rate for the hypergraph model
%R = max(delta,beta);
%beta = beta/R;
%delta = delta/R;


i_0 = [0.5]; % initial proportions of infectious individuals
T = 15; % finite time
dt = 0.1; % small time steps
n_steps = floor(T/dt); % number of time steps
%T = n_steps*dt;

%% generate hypergraph
[I,sizes_dist] = generate_hypergraph_bis(n,K);
% I is the incidence matrix of the hypergraph
%W=zeros(n,n);
n_edges=sizes_dist;
%m=n_edges(1);
fun={'2log'};
%fun={'atan'};
if strcmp(fun,'2log')
I_k = zeros(n,size(I,2));
m=n_edges(1);
I_k(:,(1:m))=I(:,(1:m));
W=(I_k)*(I_k)';
for i=2:K-1
    I_k = zeros(n,m);
    I_k(:,(m+1:m+n_edges(i)))=I(:,(m+1:m+n_edges(i)));
    W=W+2*(I_k)*(I_k)';% for f = 2*log(1+x)
 %   W=W+I(:,(m+1:m+n_edges(i+1)))*I(:,(m+1:m+n_edges(i+1)))';
    m=m+n_edges(i);
end
else
W=I*I';
end
lambda=eigs(W,1);
beta_critical = delta/(lambda);


figure;
hold on;



        
        
        %% generate initial infectious population
        l=floor(i_0*n);
        I_0 = sort(randsample(1:n,l));
        I_initial = zeros(n,1);
        I_initial(I_0,1) = ones(l,1);

        
    for b=1:length(beta)
        
        
        %% run mean field approximation model
        P = run_hypergraph_mean_field_approximation(i_0,beta(b),delta,I,T,dt,K,fun{1});
        expected_I_t = sum(P);
        subplot(1,2,1)
        if beta(b) < beta_critical
            plot(0:dt:T,expected_I_t/n,'LineWidth',3,'LineStyle','--','color','red');
        else
            plot(0:dt:T,expected_I_t/n,'LineWidth',3,'LineStyle','--','color','blue');
        end
        hold on
       
        
    end
     xlabel('Time');
     ylabel('i','Rotation',0);
     ylim([0,0.7]);
     set(gca,'FontWeight','Bold');
     set(gca,'FontSize',15);
    %% run disease for the hypergraph
    for b=1:length(beta)
        avge_I_t=zeros(1,n_steps+1);
        for i=1:runs
            I_t = zeros(n,n_steps+1); % infectious individuals evolution in the hypergraph
            I_t(:,1) = I_initial;
            S_t = ones(n,n_steps+1) - I_t; % susceptible individuals evolution in the hypergraph
            [I_t,~] = run_SIS_hypergraph_disease(beta(b),delta,I_t,S_t,I,T,dt,K,fun{1});
            avge_I_t=avge_I_t+sum(I_t);
        end
        avge_I_t=avge_I_t/runs;
        
        %% plots comparison
        
        subplot(1,2,2);
        if beta(b)<beta_critical
            plot(0:dt:T,avge_I_t/n,'LineWidth',3,'color','red');
        else
            plot(0:dt:T,avge_I_t/n,'LineWidth',3,'color','blue');
        end
        hold on
      
       
        

    end
        xlabel('Time');
        ylabel('i','Rotation',0);
        ylim([0,0.7]);
        set(gca,'FontWeight','Bold');
        set(gca,'FontSize',15);
%lgnd=legend('show');
%lgnd.Position(1) = 0.01;
%lgnd.Position(2) = 0.4;


hold off;

% ylabel('Proportion of infectious individuals')


% save figure
%ax=gca;
%saveas(ax,[pwd '/figures/alt_f_log_n_' num2str(n) '_m_' num2str(m) '_K_' num2str(K) '_T_' num2str(T) '_dt_' num2str(dt) '.png']);%,'Resolution',300);
%exportgraphics(ax,[pwd '/figures/f_min_n_' num2str(n) '_m_' num2str(m) '_K_' num2str(K) '_T_' num2str(T) '_dt_' num2str(dt) '.png'],'Resolution',300);