%% Coronoa spread model by Prof. Yinon Ashkenazy, Dr. Michael Assaf and Prof. Doron Gazit  
%% Implemented by Prof. Nadav Katz, Yuval Zamir, 
%% Samuel Goldstein and Shimon Nowik together with Roee Grant and Snir Avraham.
%% Last update: 12.4.2020 21:00
%% Contact: Prof. Nadav Katz +972-54-2563212, Yuval Zamir +972-50-8896842

%% Notations
% S - Susceptibles 
% I - Infected (IA - recovering population (hidden), 
%       ISA - recovering (known/detected) ISB - to be hospitalized)
% H - Hospitalized (HA - recovering population, HB - to become critical)
% C - Critical (CA - recovering population, CB - to die)
% D - Dead
% R - Recovered
%
% eta - DetectedI
% alpha- rI
% beta - rIH
% lambda - rIR
% delta - rHC
% delta' - rHR
% epsilon - rCD
% epsilon' - rCR

clear, close all

%% Parameters

% Time step [days]
dt = 0.1; 
% Length of simulation [days]
simLength = 250;
% Simulated time parameter
time =  dt:dt:simLength;
% Case simulated. Options are - BAU (business as usual), LockDown
% (permanent lockdown), Exit (age depended exit time) and 
% GradualExit (see documentation).
Case = 'LockDown';


%% Constant
% Patients on Day0
n0 = 10;
% Percentage of Tested out of Infected
DetectedI = 1 / 10; %eta
% Size of population
N = 9e6; 
% Percentage of population per age group
pAge = [19.7, 16.4, 14, 13.1, 11.8, 9.1, 8.1, 4.8, 3] * 1e-2;
% Numebr of age groups
nAge = length(pAge);
% Number of young age groups (e.g. yGr=6 -> 60 years differs between young and old in terms of Infection Ratio under lockdown)
yGr_LD = 6; 
% Value of Infection Rate with minimal restrictions
baserI = 1 / 3.4;
% Imposing restriction (10% shift) - based on Israel's policies
restDay0 = 40;
% Maximal number of criticl patients that can be supported by hospitals
maxCriticalPatients = 1000;

%% Rates
% Rate of Susceptabiles (S) to become Infected (I) - alpha
rI = baserI * ones(simLength/dt,nAge);
% Rate of Infected (I) to be Hospitalized (H) - beta
rIH = 1 / 5;
% Rate of Infected (I) to become Recovered (R) - lambda
rIR = 1 / 14;
% Rate of Hospitalized (H) to become Critical (C) - delta
rHC = 1 / 10;
% Rate Hospitalized (H) to Recover (R) - delta'
rHR = 1 / 20;
% Rate of Critical (C) to become Dead (D)  - epsilon
rCD = 1 / 8;
% Rate of Critical (C) to Recover (R) - epsilon'
rCR = 1 / 14;

%% Transformation Probabilities
% Probability of Detected Infected to become Hospitalized (A vector)
pIH = [0.1, 0.3, 1.2, 3.2, 4.9, 10.2, 16.6, 24.3, 27.3] * 1e-2; 
% Probability of Hospitalized to become Critical (B vector)
pHC = [5, 5, 5, 5, 6.3, 12.2, 27.4, 43.2, 70.9] * 1e-2; 
% Probability of Critical to become Dead (Gamma vector)
pCD = [40, 40, 50, 50, 48.6, 48.2, 48.4, 48.6, 48.1] * 1e-2;

%% Initialization
S   = zeros(simLength/dt,nAge); 
S(1,:) = N * pAge; % All population is susceptaible
I   = zeros(simLength/dt,nAge);
IA  = zeros(simLength/dt,nAge);
ISA = zeros(simLength/dt,nAge);
ISB = zeros(simLength/dt,nAge);
H   = zeros(simLength/dt,nAge);
HA  = zeros(simLength/dt,nAge);
HB  = zeros(simLength/dt,nAge);
C   = zeros(simLength/dt,nAge);
CA  = zeros(simLength/dt,nAge);
CB  = zeros(simLength/dt,nAge);
D   = zeros(simLength/dt,nAge);
RD  = zeros(simLength/dt,nAge);
RN  = zeros(simLength/dt,nAge);
R   = zeros(simLength/dt,nAge);

dS   = zeros(simLength/dt,nAge);
dIA  = zeros(simLength/dt,nAge);
dISA = zeros(simLength/dt,nAge);
dISB = zeros(simLength/dt,nAge);

dHA = zeros(simLength/dt,nAge);
dHB = zeros(simLength/dt,nAge);
dCA = zeros(simLength/dt,nAge);
dCB = zeros(simLength/dt,nAge);
dD  = zeros(simLength/dt,nAge);
dRD = zeros(simLength/dt,nAge);
dRN = zeros(simLength/dt,nAge);

%% Initial Conditions
I(1,:)   = n0 * pAge; 
ISA(1,:) = I(1,:).*(1-pIH)*DetectedI;
ISB(1,:) = I(1,:).*pIH*DetectedI;
IA(1,:)  = I(1,:).*(1-DetectedI);

%% Cases

if (strcmp(Case,'LockDown'))||(strcmp(Case,'Exit'))||(strcmp(Case,'GradualExit')),
    % The proccess of applying the restrictions on the public is simulated as
    % tanh function (with a time dependence) on the .
    % Using a tanh(x) sigmoid requires a rescaling to achieve 10% to 90% change
    % in restLength time, this requires a scaling factor, called  sigFactor
    sigFactor = 2.2;
    % Amount of days since the first restriction to complete lockdown 
    % based on actual data
    restLength = 12;

    % Opening days of restriction (lockdown):
    OpenDay1=restDay0+42; % for young age groups
    OpenDay2=restDay0+77; % for old age groups;
    % Intensity of quarantine (comapring to base rI):
    youngFactor = 0.4; % for young age groups
    oldFactor = 0.25; % for old age groups;
    
    minrIratio = [ones(1,yGr_LD)*youngFactor ones(1,nAge-yGr_LD)*oldFactor];
    rImin=baserI*minrIratio;

    for jj=1:nAge
        rI(:,jj)=rImin(jj)+(baserI -rImin(jj)).*(1-tanh(sigFactor*(time-(restDay0+restLength/2))/restLength))/2;  %[1/day]
    end
end

if strcmp(Case, "Exit")
    rI(OpenDay1/dt:end, 1:yGr_LD) = baserI;
    rI(OpenDay1/dt:end, yGr_LD+1:end) = baserI;
end

if strcmp(Case, "GradualExit")
    % Stage model - see documentation
    secStagerI = [0.9, 0.9, 0.9, 0.9, 0.9, 0.4, 0.25, 0.25, 0.25] * baserI;
    thirdStagerI = [1.15, 1.15, 1.15, 1.15, 1.15, 0.9, 0.25, 0.25, 0.25] * baserI;
    fourthStagerI = 1.15*ones(1,9) * baserI;
    stageModelDay0 = 89; % first day of stage stategy - 19th of April
    stageLength = 21; % length of each stage
    rI(stageModelDay0/dt:(stageModelDay0 + stageLength)/dt - 1, :) = repmat(secStagerI, stageLength/dt ,1);
    rI((stageModelDay0 + stageLength)/dt : (stageModelDay0 + 2*stageLength)/dt - 1, :) = repmat(thirdStagerI, stageLength/dt ,1);
    rI((stageModelDay0 + 2*stageLength)/dt : end, :) = repmat(fourthStagerI, length(rI((stageModelDay0 + 2*stageLength)/dt : end, :)) ,1);
end


%% Dynamics
for tt = round(time(1,1:end-1)/dt)
    
    dS(tt,:) = -rI(tt,:).*S(tt,:)*sum(I(tt,:))/N;
    dIA(tt,:) = rI(tt,:).*(1-DetectedI).*S(tt,:)*sum(I(tt,:))/N - rIR*IA(tt,:);
    dISA(tt,:) = rI(tt,:).*DetectedI.*S(tt,:).*(1-pIH)*sum(I(tt,:))/N - rIR*ISA(tt,:);
    dISB(tt,:) = rI(tt,:).*DetectedI.*S(tt,:).*pIH*sum(I(tt,:))/N - rIH*ISB(tt,:);
    dHA(tt,:) = rIH*ISB(tt,:).*(1-pHC) - rHR*HA(tt,:);
    dHB(tt,:) = rIH*ISB(tt,:).*pHC - rHC*HB(tt,:);
    if sum(C(tt,:),2) > maxCriticalPatients
        % If the number of current critical patientes is larger than the
        % capacity, all new critical patients will die.
        pCDtemp = 1; 
    else
        pCDtemp = pCD;
    end 
    dCA(tt,:) = rHC*HB(tt,:).*(1-pCDtemp) - rCR*CA(tt,:);
    dCB(tt,:) = rHC*HB(tt,:).*pCDtemp -rCD*CB(tt,:);
    dD(tt,:) = rCD*CB(tt,:);
    dRD(tt,:) = rHR*HA(tt,:) + rCR*CA(tt,:)+rIR*ISA(tt,:);
    dRN(tt,:) = rIR*IA(tt,:);
    
    S(tt+1,:)   = min(max(S(tt,:)   + dt*dS(tt,:),0),   S(1,:));
    IA(tt+1,:)  = min(max(IA(tt,:)  + dt*dIA(tt,:),0),  S(1,:));
    ISA(tt+1,:) = min(max(ISA(tt,:) + dt*dISA(tt,:),0), S(1,:));
    ISB(tt+1,:) = min(max(ISB(tt,:) + dt*dISB(tt,:),0), S(1,:));
    HA(tt+1,:)  = min(max(HA(tt,:)  + dt*dHA(tt,:),0),  S(1,:));
    HB(tt+1,:)  = min(max(HB(tt,:)  + dt*dHB(tt,:),0),  S(1,:));
    CA(tt+1,:)  = min(max(CA(tt,:)  + dt*dCA(tt,:),0),  S(1,:));
    CB(tt+1,:)  = min(max(CB(tt,:)  + dt*dCB(tt,:),0),  S(1,:));    
    D(tt+1,:)   = min(max(D(tt,:)   + dt*dD(tt,:),0),   S(1,:));
    RD(tt+1,:)  = min(max(RD(tt,:)  + dt*dRD(tt,:),0),  S(1,:));
    RN(tt+1,:)  = min(max(RN(tt,:)  + dt*dRN(tt,:),0),  S(1,:));
    
    I(tt+1,:)=IA(tt+1,:)+ISA(tt+1,:)+ISB(tt+1,:);
    H(tt+1,:)=HA(tt+1,:)+HB(tt+1,:);
    C(tt+1,:)=CA(tt+1,:)+CB(tt+1,:);
    R(tt+1,:)=RD(tt+1,:)+RN(tt+1,:);
end

%% Plots
%Adding jetColors according to age groups
values=pAge; Colors = jet; NoColors = length(Colors);
Ireduced = (values-min(values))/(max(values)-min(values))*(NoColors-1)+1;
RGB = interp1(1:NoColors,Colors,Ireduced);
set(0,'DefaultAxesColorOrder',RGB); close;

%% Plot by Age group
figure(1)
subplot(2,2,1);
plot(time,I);
xlabel('Time [days]')
ylabel('Diagnosed')
title('Diagnosed per Age Group');
legend('0-10','10-20','20-30','30-40','40-50','50-60','60-70','70-80','80+')

subplot(2,2,2);
plot(time,H);
xlabel('Time [days]')
ylabel('Hospitalized')
title('Hospitalized ');

subplot(2,2,3);
plot(time,C);
xlabel('Time [days]')
ylabel('Critical Patients')
title('Critical Patients ');

subplot(2,2,4);
plot(time,D);
xlabel('Time [days]')
ylabel(' ')
title('Dead Patients ');

for jj = 1:4 %Num subplots per figure
    subplot(2,2,jj);
    set(gca,'FontSize',12);
end


%% Plot total
s = figure(2);
subplot(2,2,1);
sumI = sum(I,2);
sPlot = plot(time,sumI);

xlabel('Time [days]')
ylabel('Total Diagnosed')
title('Total Diagnosed ');
[maxI, timemaxI] = max(sumI);
cursorMode = datacursormode(s);
sDatatip = cursorMode.createDatatip(sPlot);
pos =  [timemaxI * dt maxI 0];
set(sDatatip, 'Position', pos) 
updateDataCursors(cursorMode);


subplot(2,2,2);
sumH = sum(H,2);
sPlot = plot(time,sumH);
xlabel('Time [days]')
ylabel('Total Hospitalized')
title('Hospitalized ');
[maxH, index] = max(sumH);
cursorMode = datacursormode(s);
sDatatip = cursorMode.createDatatip(sPlot);
pos =  [index * dt maxH 0];
set(sDatatip, 'Position', pos) 
cursorMode.UpdateFcn = @displayCoordinates;
updateDataCursors(cursorMode);



subplot(2,2,3);
sumC = sum(C,2);
sPlot = plot(time, sumC);
xlabel('Time [days]')
ylabel('Critical Patients')
title('Critical Patients ');
[maxC, index] = max(sumC);
cursorMode = datacursormode(s);
sDatatip = cursorMode.createDatatip(sPlot);
pos =  [index * dt maxC 0];
set(sDatatip, 'Position', pos) 
cursorMode.UpdateFcn = @displayCoordinates;
updateDataCursors(cursorMode);


subplot(2,2,4);
sumD = sum(D,2);
sPlot = plot(time, sumD);
xlabel('Time [days]')
ylabel('Total Deaths')
title('Dead Patients ');
[maxD, index] = max(sumD);
cursorMode = datacursormode(s);
sDatatip = cursorMode.createDatatip(sPlot);
pos =  [index * dt maxD 0];
set(sDatatip, 'Position', pos) 
cursorMode.UpdateFcn = @displayCoordinates;
updateDataCursors(cursorMode);

for jj = 1:4 %Num subplots per figure
    subplot(2,2,jj);
    set(gca,'FontSize',12);
    ylim([1 inf])
end



%% Plot log scale
s = figure(3);
subplot(2,2,1);
sPlot = semilogy(time,sum(I,2),'LineWidth',2);
xlabel('Time [days]')
ylabel('Total Infected')
title('Total Infected ');

subplot(2,2,2);
sPlot = semilogy(time,sum(H,2),'LineWidth',2);
xlabel('Time [days]')
ylabel('Total Hospitalized')
title('Hospitalized ');

subplot(2,2,3);
sumC = sum(C,2);
sPlot = semilogy(time, sumC,'LineWidth',2);
xlabel('Time [days]')
ylabel('Critical Patients')
title('Critical Patients ');

subplot(2,2,4);
sPlot = semilogy(time,sum(D,2),'LineWidth',2);
xlabel('Time [days]')
ylabel('Total Deaths')
title('Dead Patients ');

for jj = 1:4 %Num subplots per figure
    subplot(2,2,jj);
    set(gca,'FontSize',12);
    ylim([10^(0) inf])
end

% Other Parameters
figure(999)
total_people = S + I + H + D + C + R;
subplot(2,1,1)
plot(time, sum(total_people,2)/N)
title('Population Conservation (Ratio)')
xlabel('Time')
ylabel('People at T / People at T0')
subplot(2,1,2);
plot(time,rI);
title('Infection Rate');
xlabel('Time');
ylabel('Infection rRate');
legend('0-10','10-20','20-30','30-40','40-50','50-60','60-70','70-80','80+') 

figure(3)
figure(2)
figure(1)

function txt = displayCoordinates(~,info)
    x = info.Position(1);
    y = round(info.Position(2),3,'significant');
    txt = ['Day: ' num2str(x)  newline 'Max Value: ' num2str(y,'%.2e')];
end