% To generate heatmaps in Figure 4, the probability of developing drug-resistance (qtau)
% should be selcted as indicated in the main simulation part of the code
% below. For each qtau, the code will generate subplots of the total number of infections 
% (epidemic final size) as well as the cumulative number of drug-resistant infections.  

clear all
clc
%% System parameters with no demographics
mu=0;
gamma=1/(5);
R0=3;
beta=R0*(mu+gamma);

%% Variables and their dimension
h = 0.05;                   % simulation time-step
tn = 3*356;                 % total simulation time in days
t = (0:h:tn);               % time vector    
nt = length(t);             % length of the time vector
S = zeros(nt,1);            % vector of the susceptible population
I = zeros(nt,1);            % vector of the infected population with the wild-type infection
Ir = zeros(nt,1);           % vector of the infected population with the drug-resistant infection
T = zeros(nt,1);            % vector of the effectively treated infection 
Totalr = zeros(nt,1);       % vector of cumulative number of drug-resitant infections 

%% Initial conditions       
Ir0 = 0.0;              % initial fraction of population with the drug-resistance infection
I0 = 0.001;             % initial fraction of population with the wild-type infection
S0 = 1-I0-Ir0;          % initial fraction of population susceptible
T0 = 1-I0-Ir0-S0;       % initial fraction of population effectively treated
Totalr0 = 0;            % initial cumulative number of drug-resistant-infections

%% vector of delay in start of treatment (tau)
dtau=19;                % number of points for dicretizing 'tau'
h1=(1/gamma-0.25)/dtau; % increment of 0.25 days for the delay vector between 0.25 and 1/gamma = 5 days 

%% vector of drug-resistance fitness
Ndelta=19;             % number of points for discretizing 'delta'
%TTA = zeros(dtau+1,1);
%FTA = zeros(dtau+1,1); 

%% initializing vectors for outputs
S(1) = S0;
I(1) = I0;
Ir(1) = Ir0;
T(1) = T0;
Totalr(1) = Totalr0;
TTA = zeros(dtau+1,Ndelta+1);   % vector of final sizr for healthmaps
FTA = zeros(dtau+1,Ndelta+1);   % vector of cumulative number of drug-resistance for heatmaps

%% main simulations
 for kk=1:Ndelta+1                  % loop for drug-resistance fitness
     ddelta(kk)=(kk-1)/Ndelta;      % vector of fitness
     delta=ddelta(kk);              % drug-resistance fitness
    for ii=1:dtau+1                     % loop for delay in start of treatment
        ttau(ii)=0.25+(ii-1)*h1;        % vector of delay
        tau=ttau(ii);                   % delay in start of treatment
        %% functional forms of probability of developing drug-resistance with delay in start of treatment
        qtau = 15*exp(-2*tau)/(1 + exp(-2*(tau-2)));        % this qtau give Fig4 (A,D)       
        %qtau = 2*exp(-tau)/(1 + exp(-5*(tau-1.5)));        % this qtau give Fig4 (B,E)
        %qtau = 0.25/(1 + exp(-2*(tau-2)));                 % this qtau give Fig4 (C,F)
        %% main simulations with time
        delay=tau/h;
        for j = 2:nt
            S(j) = (S(j-1)+h*mu)/(1+h*beta*(I(j-1)+delta*Ir(j-1))+h*mu);
            if j<delay+1
                I(j) = (I(j-1)+h*beta*S(j-1)*I(j-1))/(1+h*gamma+h*mu);
                Ir(j) = (Ir(j-1)+h*delta*beta*S(j-1)*Ir(j-1))/(1+h*gamma+h*mu);
                T(j) = (T(j-1)+h*gamma*(I(j-1)+Ir(j-1)))/(1+h*mu);
                Totalr(j) = Totalr(j-1) + h*gamma*Ir(j-1);
            else
                I(j) = (I(j-1) + h*beta*S(j-1)*I(j-1) - h*beta*S(j-delay)*I(j-delay)*exp(-(mu+gamma)*tau))/(1 + h*gamma + h*mu);
                Ir(j) = (Ir(j-1) + h*delta*beta*S(j-1)*Ir(j-1) + h*qtau*beta*S(j-delay)*I(j-delay)*exp(-(mu+gamma)*tau))/(1 + h*gamma + h*mu);
                T(j) = (T(j-1) + h*gamma*(I(j-1)+Ir(j-1)) + h*(1-qtau)*beta*S(j-delay)*I(j-delay)*exp(-(mu+gamma)*tau))/(1 + h*mu);
                Totalr(j) = Totalr(j-1) + h*gamma*Ir(j-1);
             end
         end
TTA(ii,kk)=S0-S(nt);        % total number of infections                        
FTA(ii,kk)=Totalr(nt);      % total number of drug-resiatant infections
%FTA(ii)=(S0-S(nt))*exp(-gamma*tau);        % fraction of infected population treated.      
     end
 end

 %% plotting heatmaps
 subplot(2,1,1)
 M1=interpn(TTA,7);
 imagesc(ttau,ddelta,M1');
 set(gca,'Ydir','Normal');
 colormap(jet(10));
 
 subplot(2,1,2)
 M2=interpn(FTA,7);
 imagesc(ttau,ddelta,M2');
 set(gca,'Ydir','Normal');
 colormap(jet(10));