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.05, 0.3, 0.7]; % 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.05, 0.95]; % initial proportions of infectious individuals
T = 40; % finite time
dt = 0.05; % small time steps
n_steps = floor(T/dt); % number of time steps
%T = n_steps*dt;

%% generate hypergraph
[I,n_edges] = generate_hypergraph_bis(n,K);

% I is the incidence matrix of the hypergraph

fun = {'atan','log','min','collective','2log'};
for j=4:4
    if j==4
        beta = [0.1, 0.5, 1];
    end
figure;
hold on;


for p=1:2
        
        
        %% generate initial infectious population
        l=floor(i_0(p)*n);
        I_0 = sort(randsample(1:n,l));
        I_initial = zeros(n,1);
        I_initial(I_0,1) = ones(l,1);
        
        
        
        
    for b=1:3
        
        avge_I_t=zeros(1,n_steps+1);
        %% run mean field approximation model
        P = run_hypergraph_mean_field_approximation(i_0(p),beta(b),delta,I,T,dt,K,fun{j});
        expected_I_t = sum(P);
        %% run disease for the hypergraph
        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{j});
            avge_I_t=avge_I_t+sum(I_t);
        end
        avge_I_t=avge_I_t/runs;
        
        %% plots comparison
        %subplot(3,2,p+2*(b-1))
        subplot(2,3,b+3*(p-1));
        %if (p==2 && b ==3)
        %    plot(0:dt:T,I_t/n,'LineWidth',3,'DisplayName', ['exact' newline 'model']);
        %    hold on
        %    plot(0:dt:T,expected_I_t/n,'LineWidth',3,'DisplayName',['mean field' newline 'approxi-' newline 'mation']);
        %else
        plot(0:dt:T,avge_I_t/n,'LineWidth',3);
        hold on 
        plot(0:dt:T,expected_I_t/n,'LineWidth',3,'LineStyle','--');
       
        %end
       
        %p2=plot(0:dt:T,expected_I_t/n);
        %legend('exact model','mean field approximation');
        title(['i_0 = ' num2str(i_0(p)) ', \beta = ' num2str(beta(b))]);
        xlabel('Time');
        % customize
        %set(p1,'LineWidth',3);
        %set(p2,'LineWidth',3)
        set(gca,'FontWeight','Bold');
        set(gca,'FontSize',15);

    end
end
%lgnd=legend('show');
%lgnd.Position(1) = 0.01;
%lgnd.Position(2) = 0.4;


hold off;

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