%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   
%   SISAveraging
%
%   In this code we will numerically an SIS epidemic on a random network
%   and average over the dynamics and and compare with the ODE solution. 
%   
%   In the network each node will have 
%   a status of either S,I denoting if the node is susceptible, or infected.
%   We will also keep track of the edge stats SS, SI, and II. 
%   At each time step susceptible nodes will have a probability of becoming 
%   infected and infected nodes will have a probability of recovering back 
%   to susceptible.
%
%   The following rules will be implemented to change the status of nodes
%   at each time step
%
%   1. Nodes that are susceptible will become infected with probability:
%       p_{S\rightarrow I}=\beta \times dt \times # of infected
%       neighbors
%
%   2. Nodes that are infected will recover with probability:
%       p_{I\rightarrow R}=\alpha \times dt
%
%   The code will be split into three sections. The first section will
%   generate the adjacency matrix for the system and initilize the vectors
%   which encode the status of each node. The second section will loop
%   through the simulation applying the rules above. The third section will
%   solve the corresponding differential equation. Note, all probabilities 
%   and transitions will be implemented simultaneously.
%
%   Inputs:
%   1. S0= initial number of susceptible individuals
%   2. I0= initial number of infected individuals
%   3. beta= force of infection
%   4. alpha= recovery rate
%   5. Tmax= total time of the simulation
%   6. n= Number of temporal points to sample over.
%   7. p= probability for generating connections in network
%   8. M=number of simulations
%
%   Outputs:
%   Plot of average dynamics over M trials compared with ODE model
%
%   Movie of the simulation.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Inputs

S0=999;
I0=1;

beta=.2;
alpha=.5;

Tmax=4;

n=1000;

p=.05;

M=10;

%% Derived quantites
% Here is where I put quantities that are computed from the inputs, for
% instance dt.

dt=Tmax/(n-1);

%% Construction of Initial States

%First we just compute the total number of nodes needed:
N=S0+I0;

%We next create matrices that will track the status the nodes with respect
%to time. In these matrices the row number will correspond to the node
%while the column will correspond to the time.

S=zeros(N,n);
I=zeros(N,n);

%We now fill in the first columns of S and I with there initial states
S(1:S0,1)=ones(S0,1);
I(N-I0+1:end,1)=ones(I0,1);

%The matrices S and I encode the state of the disease at each time. When we
%plot the state of the disease we will need to look at each column and have
%a way of extracting which nodes are susceptible or infected. We will need
%do this by tagging the nodes by there index. We do this by creating a
%a vector of the form Tag=[1,2,3,...,N]. We will see how to use this when we
%plot the graph with nodes indicated by their infection status.

Tag=linspace(1,N,N)';

%% Construction of Adjacency Matrix
% We next generate the adjacency matrix. Each node i will with probabiliy p
% be connected with node j. We do this by creating an upper triangular
% matrix with 1's and 0's with a 1 indicating that node i is connected to
% node j. Since the adjacency matrix is symmetric we just add the upper
% tringular matrix to its transpose to give the full symmetric adjacency
% matrix.


A = rand(N,N) < p;
A = triu(A,1);
A = A + A';

k=1/N*sum(sum(A));


%% Simulation of network dynamics

%The following command makes sure we do not keep making new figures
%everytime we plot.

Savg=zeros(n,1);
    Iavg=zeros(n,1);
    SSavg=zeros(n,1);
    SIavg=zeros(n,1);
    IIavg=zeros(n,1);

for j=1:M

    
    
    S(1:S0,1)=ones(S0,1);
    I(N-I0+1:end,1)=ones(I0,1);
for i=1:n-1
    %First I want to generate a vector of probablities for each node. I
    %will use this vector to decide what each node should do in the next
    %time step.
    
    VecProb=rand(N,1);
    
    %Now I want to calculate how many infections there are at an arbitrary
    %node j. This can be done by matrix multiplication with the adjacency
    %matrix since InfectionEdges=A*I(:,i) is a vector with the j-th row
    %corresponding to the number of infections connected to the j-th node.
    
    NumInfections=A*I(:,i);
    
    %We now need to check if a susceptible become infected and assign a new
    %vector to the individuals that will become infected. We do this by
    %checking which of the current susceptibles are connected to an infected
    %individual and then check if they become infected using our
    %probability rules.
    
    ToInfect=VecProb<NumInfections.*S(:,i)*beta*dt;
    
    %We now check to see which infected individuals recover. This does not
    %depend on the neighbors so the computation does not require the
    %adjacency matrix:
    
    ToRecover=VecProb<I(:,i)*alpha*dt;
    
    %Next we update the columns of S and I to reflect there changes in
    %status.
    
    S(:,i+1)=S(:,i)-ToInfect+ToRecover;
    I(:,i+1)=I(:,i)+ToInfect-ToRecover;
    
    
    
  
end
STot=sum(S);
ITot=sum(I);

SS=diag(S'*A*S);
SI=diag(S'*A*I);
II=diag(I'*A*I);

SS0=SS(1);
SI0=SI(1);
II0=II(1);

Savg=Savg+STot';
Iavg=Iavg+ITot';
SSavg=SSavg+SS;
SIavg=SIavg+SI;
IIavg=IIavg+II;

end

Savg=Savg/M;
Iavg=Iavg/M;
SSavg=SSavg/M;
SIavg=SIavg/M;
IIavg=IIavg/M;

%% Plot of susceptible, infected, and recovered and edges as functions of time

%Domain to plot over
t=linspace(0,Tmax,n);

%Computing the sum of each column to get the total number of S,I,R for each
%time step



%% Construction of Temporal Domain for the simulation

%Creation of a vector of n equally spaced points from 0 to Tmax
tspan=linspace(0,Tmax,n); 


%% Simulation of ODE45
% ODE45 takes in four things in order to simulate a system of ODES:
%   a) @(t,y) representing the independent and depenendant variables. In
%   this case y is a row vector with three components y=(S,I,R). This part
%   of the code is just telling ODE45 which are the variables that are
%   going to change in time. 
%
%   b) A function which represents the righthand side of the ODEs and will
%   be defined later in the code.
%
%   c) tspan which is the set of temporal points we are sampling over
%
%   d) y0 which is a vector of initial conditions y0=(S0,I0,R0)^T. 

% Form initial conditions into a column vector, ' denotes transpose
y0=[S0,I0,SS0,SI0,II0]'; 

%The Implementation of ODE45. Note, odefcn is the righthand side for the
%system of ODEs and is defined at the end of the code. 

[t,y] = ode45(@(t,y) odefcn(t,y,k,beta,alpha,N), tspan, y0);

%% Extracting output and plotting the results
SODE=y(:,1);
IODE=y(:,2);
SSODE=y(:,3);
SIODE=y(:,4);
IIODE=y(:,5);


%Basic plot command, here I am plotting three curves and setting all of
%them to have length 2. 

%Plotting everything
figure('units','normalized','outerposition',[0 0 1 1]); 
plot(t,Savg,t,Iavg,t,SODE,t,IODE,'linewidth',2);

xlabel('$$t$$','interpreter','latex')
ylabel('Population','interpreter','latex')

ax=gca;
ax.FontSize = 40;

set(gcf,'color','w');

legend('$$S$$ Average','$$I$$ Average','$$S$$ ODE','$$I$$ ODE','interpreter','latex','fontsize',40)

%Labeling axes. Note, Matlab can do LaTeX!
xlabel('$$t$$','interpreter','latex')
ylabel('Population','interpreter','latex')

%Setting the fontsize of the labels and ticks. Note, when you make your
%final figures you might have to shrink or enlarge your fontize so that
%your figures have readible text that matches the font size in the paper. 
ax=gca;
ax.FontSize = 40;

%This code just sets a white background.
set(gcf,'color','w');


figure('units','normalized','outerposition',[0 0 1 1]); 
plot(t,SSavg,t,SIavg,t,IIavg,t,SSODE, t, SIODE, t, IIODE,'linewidth',2);

%Labeling axes. Note, Matlab can do LaTeX!
xlabel('$$t$$','interpreter','latex')
ylabel('Edges','interpreter','latex')

%Setting the fontsize of the labels and ticks. Note, when you make your
%final figures you might have to shrink or enlarge your fontize so that
%your figures have readible text that matches the font size in the paper. 
ax=gca;
ax.FontSize = 40;

%This code just sets a white background.
set(gcf,'color','w');

%Inserting a legend since we are plotting three curves. 
legend('$$[SS]$$ Average','$$[SI]$$ Average','$$[II]$$ Average','$$[SS]$$ ODE', '$$[SI]$$ ODE', '$$[II]$$ ODE','interpreter','latex','fontsize',40)


%% Function which represents the right hand side of the odes. 
% This is the function that is passed to ODE45. 
%
%   Inputs
%   a) t the current timeplot
%
%   b) y representing a column of (S,I,R) at time t. 
%
%   c) beta, alpha are the constant parameters
%
%   Outpus
%   a) dy/dt which is a row vector with components
%   (\dot{S},\dot{I},\dot{R}) denoting the values of the time derivative
%   sof S,I,R respectivly. 

function dydt = odefcn(t,y,k,beta,alpha,N)

%Conversion of y into the variables we are using for the problem. I do this
%to make the debugging easier and to make the code more readable. 

S=y(1);
I=y(2);
SS=y(3);
SI=y(4);
II=y(5);

SSI=(k-1)/k*SS*SI/S;
ISI=(k-1)/k*SI^2/S;


Sdot=-beta*SI+alpha*I;
Idot=beta*SI-alpha*I;
SSdot=-2*beta*SSI+2*alpha*SI;
SIdot=beta*SSI+alpha*II-beta*ISI-alpha*SI-beta*SI;
IIdot=2*beta*SI+2*beta*ISI-2*alpha*II;

%Putting the derivatives into a single row vector dydt.
dydt(1)=Sdot;
dydt(2)=Idot;
dydt(3)=SSdot;
dydt(4)=SIdot;
dydt(5)=IIdot;

dydt=dydt';


end

