clear

plotstuff=1;
stochasticpaths=0;
exportimages=0;
calcerrors=0;

% Ranges for plotting images
if plotstuff
    R0range=[1.5 2 6];
    Nrange=[200];
    ICrange=[1 2 5 10];
    ICrange=[1 5 10];
end
% Ranges for calculating "errors" between ODE and Kolmogorov eq
if calcerrors
    R0range=[1.01:0.5:15];
    Nrange=[200];
    ICrange=[1:15];
end

for indR0=1:length(R0range)           % loop over R0 values
    for indN=1:length(Nrange)         % loop over total population size
        for indIC=1:length(ICrange)   % loop over initial number of infected individuals
            R0=R0range(indR0)
            N=Nrange(indN);
            IC=ICrange(indIC)
            
            b=R0/N; % nondimensional infection rate constant (beta/gamma)
            g=1;    % nondimensional recovery rate constant - leave as 1 (time is scaled by 1/gammma)
            
            MinNondimTimeScale=min(1/g,1/(R0*g))
            TimeSteps=round(1000*MinNondimTimeScale);
            dt=MinNondimTimeScale/10;  % solution discretization chosen so as to be 1/10 of the fastest time scale
            
            % Solve the Kolmogorov equations
            A=zeros(N,N); % initialize the ODE matrix
            
            % Fill in the ODE matrix
            for n=1:N
                if n==1
                    A(n,n:n+1)=[-b*n*(N-n)-g*n g*(n+1)];
                elseif n==N
                    A(n,n-1:n)=[b*(n-1)*(N-n+1) -b*n*(N-n)-g*n];
                else
                    A(n,n-1:1:n+1)=[b*(n-1)*(N-n+1) -b*n*(N-n)-g*n g*(n+1)];
                end
            end
            
            pIC=zeros(N,1);          % initialize the initial condition vector (all zero probs)
            pIC(IC)=1;               % set the initial probability of IC infected individuals to be 1
            time=[1:TimeSteps]*dt;   % vector for plotting functions of time
            poft=zeros(TimeSteps,N); % vector of probabilities (the solution)
            
            for m=1:TimeSteps
                poft(m,:)=expm(A*m*dt)*pIC; % analytical solution to the matrix ODE evaluated at discretization
            end
            poft2=[[1-sum(poft,2)] poft];  % add in a row for p0(t)=1-sum(pn(t))
            
            yoft=[0:N]*poft2';  % the expected value / mean number of infected individuals
            soft=sqrt([0:N].^2*poft2'-yoft.^2); % standard deviation
            
            % Solution to the associated logistic ODE:
            K=(N-g/b);
            C=IC/(K-IC);
            r=b*N-g;
            Ioft=K*C*exp(r*time)./(1+C*exp(r*time));
            
            difference(indR0,indIC)=sum(abs(Ioft-yoft)); % a measure of the difference between the mean and the ODE solution

            if plotstuff
                figure(1)
                clf
                hold on

                if stochasticpaths      % toggle that lets me turn this section on/off
                    % A handful of solutions to the SDE. Note that this is not the best (most efficient) way 
                    % to generate stochastic realizations but it is the simplest conceptually. 
                    % For further reading, look for "SSA Direct method", "tau-leaping"
					
                    nStoch=20;            % number of realizations
                    nExtinct=0;           % a count of extinctions
                    for kk=1:nStoch       % draw realizations from a hat...
                        reactionTimes=0;  % vector of reaction times (this sets i=1)
                        nInfected=IC;     % vector of successive state values (this sets i=1)
                        reaction=2;       % a count of reactions

                        while reactionTimes(reaction-1)<TimeSteps*dt  % go only as far in time as does poft
                            if nInfected(reaction-1)==0               % check if the infection died out
                                reactionTimes(reaction)=TimeSteps*dt;   % set a final reaction time
                                nInfected(reaction)=0;                  % set a final state 
                                nExtinct=nExtinct+1;                    % increment extinction count
                                break                                   % break out of while loop
                            end
                            Rinfect=b*nInfected(reaction-1)*(N-nInfected(reaction-1)); % update the infection rate
                            Rrecover=g*nInfected(reaction-1);                          % update the recovery rate
                            if nInfected(reaction-1)~=N               % check that there are still some susceptibles around
                                tinfect=-log(rand)/Rinfect;           % draw a next-infection-time from the hat
                                trecover=-log(rand)/Rrecover;         % draw a next-recovery-time from the hat
                                [tnext,type]=min([tinfect trecover]); % pick the earlier of these as the next reaction time and make note of the reaction type
                            else
                                tnext=-log(rand)/Rrecover;            % if everyone is infected, the next reaction is a recovery
                                type=2;                               % reaction type 1 = infection, reaction type 2 = recovery
                            end                
                            nInfected(reaction)=nInfected(reaction-1)+(type==1)-(type==2); % update the state
                            reactionTimes(reaction)=reactionTimes(reaction-1)+tnext;       % add the next-reaction time to clock time
                            reaction=reaction+1;                                           % increment reaction count
                        end

                        temp=[reactionTimes' reactionTimes']';             % reshape reaction time data for plotting purposes
                        reactionTimesPlot=zeros(length(reactionTimes),2);  % same
                        reactionTimesPlot=temp(2:end);                     % same
                        temp=[nInfected' nInfected']';                     % reshape state data for plotting purposes
                        nInfectedPlot=zeros(length(nInfected),2);          % same
                        nInfectedPlot=temp(1:end-1);                       % same
                        plot(reactionTimesPlot,nInfectedPlot,'m','linewidth',1);  % plot stochastic realization
                    end                 % loop that creates a realization
                end                     % end stochastic paths toggle
                

                h=imagesc(time,[-0.5 N-0.5],poft2');       % plot intensity image of pn(t)
                alpha(h,0.85)
       %         plot(time,Ioft,'g','linewidth',2)  % plot logistic solution

       %         plot(time,yoft,'m','linewidth',2)  % plot expected value
       %         plot(time,yoft+soft,'m--','linewidth',2)  % plot expected value + sigma
       %         plot(time,yoft-soft,'m--','linewidth',2)  % plot expected value - sigma
       %         plot(time,poft2(:,1)/max(poft2(:,1))*N/5,'y')  % plot p0(t) (scaled up by N/5 so you can see it)

                str=sprintf('Max p_0(t) = %0.2f',max(poft2(:,1)));  % not really the max but practically it is for the time window
                texthandle=text(TimeSteps*dt*5/8-dt/8,N*(1/4)-1/2,str,'Color','k','FontSize',16);
                texthandle=text(TimeSteps*dt*5/8,N*(1/4),str,'Color','y','FontSize',16);

                str=sprintf('R_0(t) = %0.1f,  N = %0.0f, IC = %0.0f',R0,N,IC);
                texthandle=text(TimeSteps*dt/20-dt/8,N*(9/10)-1/2,str,'Color','k','FontSize',16);
                texthandle=text(TimeSteps*dt/20,N*(9/10),str,'Color','y','FontSize',16);

                if stochasticpaths
                    str=sprintf('%0.0f out of %0.0f infections died out',nExtinct,nStoch);
                    texthandle=text(TimeSteps*dt/20-dt/8,N*(8/10)-1/2,str,'Color','k','FontSize',16);
                    texthandle=text(TimeSteps*dt/20,N*(8/10),str,'Color','y','FontSize',16);
                end
%                legend('Logistic solution I(t)','Stochastic x(t)','Location','northwest')
                axis([dt TimeSteps*dt 0 N])
                xlabel('Time')
                ylabel('Number of infected individuals')
                set(gca,'fontsize',16)
                clmp=[(pi/2+atan(50*([0:63]'/63-0.1)))/pi]*[1 1 1];  % modify the colormap - increase contrast for visibility
                colormap(clmp)
%                colormap gray
                
                if exportimages
                    str=sprintf('FigureR_0%0.1fN%0.0fIC%0.0f.png',R0,N,IC);
                    saveas(gcf,str)
                end
                pause(0.1)
            end
        end
    end
end

%%
if calcerrors
    figure(2)
    clf
    hold on
    imagesc(ICrange,R0range,difference)
    xlabel('Initial number of infected individuals','FontSize',16)
    ylabel('R_0','FontSize',16)
    title('\int_0^t  |I(t)-y(t)|  dt','FontSize',16)
    set(gca,'fontsize',16)
    saveas(gcf,'FigureErrorPlot.png')
end