%% Modification to Covid-19 Model
% (c) S.Bigio - based on slides from April 2021
close all; clear;

% Convergence
tol=0.001; % tolerance
greed=0.1;

% Time span in days
days=1000     ;
time=1:days;

% Paramters:
mil=10e5    ; % definition of million
N=30*10e5   ; % US pop
gamma=1/24  ; % recovery rate

% Figure
Daleth_h=0.15; % infectious rate - high
Daleth_l=Daleth_h/5; % infectious rate - low

% Initial COnditions
S_0=0.99999*N ; % Initial succeptible pop
sigma_h=0.99  ; % fraction of high succetiple (initial date)
R_0=0        ;  % initial recovered
D_0=0        ;  % initial dead
I_0=N-(S_0+R_0+D_0); % initial infected

% Switching
theta=200.5; % extreme value parameter for activity switch
ext=@(v1,v2) 1/(1+exp(theta*(v2-v1))); % extreme valued fucntion

% Preferences
beta=1-(1-0.88)/360                              ; % daily discount
u_E=100                                          ; % End of Pandemic utility flow
u_h=100                                          ; % utility flow during pandemic - high activity
u_l=50                                           ; % utility flow during pandemic - low activity
u_r=u_h                                          ; % utility flow during pandemic - recoverd 
u_i=0                                            ; % utility flow during pandemic - infected
V_D=-1000*(1-beta)/(1-beta)                      ; % value of begin dead
   
% health capacity
uci_threshold=5e-3; % fraction of UCI beds / risk of death (multiply by death risk to get effective beds)
delta=0.5    ; % death rate at critical condition without UCI care 
delta_n=delta; % death rate at critical condition without UCI care
delta_h=delta/10  ; % death rate at UCI unit

% matching function
m=@(S,I,R,D) (S*I)/(S+R+I)^2; % matching function

% Seasonal Pattern - Covid
freq=time*3*2*pi/days; % Cycle
phase_i=1/4*pi       ; % what part of the year does cycle start (2pi=one year)
scale_se=1           ; % intensity at peak of cycle
season_p=3           ; % how intense is the cycle >1 (1=sine function) 
Daleth_se=0.5*scale_se*sin(phase_i+freq).^season_p+1; % increase in seasonal pattern

% Quarantine Policy
q_se=1.0     ; % how penalizing is high activity during policy 
phase=-1*pi/6; % when does the policy start relative to seasonal cycle
pol_freq=1   ; % 1 if cycles coincide
policy_p=1   ; % how intense is the policy
u_h_se=1-q_se*0.5*((scale_se*sin(phase_i+pol_freq*freq+phase)).^policy_p+1);

% UCI growth
rate_uci=0.72; % 0.72/r where r is doubling every r years (baseline double in one year)
UCI_g=@(tt) exp(rate_uci*tt/365); 

%% Vectorization
% Vectorize
S_t=S_0*ones(days,1);
D_t=D_0*ones(days,1);
I_t=I_0*ones(days,1);
R_t=R_0*ones(days,1);
m_t=(1-S_0/(S_0+I_0+R_0))*I_0/(R_0+I_0)*ones(days,1);
f_0=m_t(1);
    
% Vectorize
S_t=S_0*ones(days,1);
S_h_t=sigma_h*S_t;
S_l_t=(1-sigma_h)*S_t;
D_t=D_0*ones(days,1);
I_t=I_0*ones(days,1);
R_t=R_0*ones(days,1);
m_t=m(S_0,I_0,R_0,D_0)*ones(days,1);
p_h_t=0*ones(days,1);
p_l_t=0*ones(days,1);
p_t=0*ones(days,1);

% Initial Guess - Everyone stays doing business as usual
phi_hh_t=ones(days,1);
phi_ll_t=zeros(days,1);

%% Post pandemic values
% Computing Values
V_R=u_E*(1-beta)/(1-beta)                                     ;
V_h_ss  = V_R;
V_l_ss  = V_R;
phi_hh_ss=0.5;
phi_ll_ss=0.5;
for ii=1:100
    phi_hh_ss = ext(V_h_ss,V_l_ss);
    phi_ll_ss = ext(V_l_ss,V_h_ss);
    for jj=1:100
        V_h_ss    = u_h*(1-beta) + beta*(phi_hh_ss*V_h_ss+(1-phi_hh_ss)*V_l_ss);
        V_l_ss    = u_l*(1-beta) + beta*(phi_ll_ss*V_l_ss+(1-phi_ll_ss)*V_h_ss);
    end
end
% Recovery, you start working    
V_R=V_h_ss;



% Compute Transition
cond=2*tol;
while cond>tol
    hh_in=phi_hh_t;
    ll_in=phi_ll_t;
    delta_t=delta_h*ones(days,1);
    for tt=time   
        Is_h=Daleth_se(tt)*Daleth_h*m_t(tt)*S_h_t(tt);
        Is_l=Daleth_se(tt)*Daleth_l*m_t(tt)*S_l_t(tt);
        HL_t=((1-phi_hh_t(tt))*S_h_t(tt)-(1-phi_ll_t(tt))*S_l_t(tt));
        S_h_t(tt+1)=S_h_t(tt)-Is_h-HL_t;
        S_l_t(tt+1)=S_l_t(tt)-Is_l+HL_t;    
        S_t(tt+1)=S_l_t(tt+1)+S_h_t(tt+1);
        I_t(tt+1)=I_t(tt)+Is_h+Is_l-gamma*I_t(tt);
        delta_t(tt+1)=delta_h*min(I_t(tt)/N,uci_threshold*UCI_g(tt))+(delta_n-delta_h)*max(I_t(tt)/N-uci_threshold*UCI_g(tt),0);
        R_t(tt+1)=R_t(tt)+gamma*(1-delta_t(tt+1))*I_t(tt);
        D_t(tt+1)=D_t(tt)+gamma*(delta_t(tt+1))*I_t(tt);
        m_t(tt+1)=m(S_t(tt+1),I_t(tt+1),R_t(tt+1),D_t(tt+1));
        p_h_t(tt+1)=Daleth_se(tt)*Daleth_h*m_t(tt+1);
        p_l_t(tt+1)=Daleth_se(tt)*Daleth_l*m_t(tt+1);
        p_t(tt+1)=(Is_h+Is_l)/S_t(tt);
    end
    % Shrinck
    S_t=S_t(1:end-1);
    S_h_t=S_h_t(1:end-1);
    S_l_t=S_l_t(1:end-1);
    D_t=D_t(1:end-1);
    I_t=I_t(1:end-1);
    R_t=R_t(1:end-1);
    p_l_t=p_l_t(1:end-1);
    p_h_t=p_h_t(1:end-1);
    p_t=p_t(1:end-1);
    m_t=m_t(1:end-1);
    delta_t=delta_t(1:end-1);

    % Computing Values
    V_I=(u_i*(1-beta)+beta*gamma*(V_R*(1-delta_t(end))+delta_t(end)*V_D))/(1-beta*(1-gamma))*ones(days,1);            
    V_h_t  = V_R*ones(days,1);
    V_l_t  = V_R*ones(days,1);
    for tt=days:-1:2
        V_I(tt-1)      = u_i*(1-beta) + beta*((1-gamma)*V_I(tt)+gamma*(delta_t(tt)*V_D+(1-delta_t(tt))*V_R));
        V_h_t(tt-1)    = u_h_se(tt)*u_h*(1-beta) + beta*((1-p_h_t(tt))*(phi_hh_t(tt)*V_h_t(tt)+(1-phi_hh_t(tt))*V_l_t(tt))+p_h_t(tt)*V_I(tt));
        V_l_t(tt-1)    = u_l*(1-beta) + beta*((1-p_l_t(tt))*(phi_ll_t(tt)*V_l_t(tt)+(1-phi_ll_t(tt))*V_h_t(tt))+p_l_t(tt)*V_I(tt));
    end

%     V_h_t  = V_R*ones(days,1);
%     V_l_t  = V_R*ones(days,1);
    phi_hh_t=p_t;
    phi_ll_t=p_t;
    for tt=days:-1:2
        phi_hh_t(tt)   = ext(V_h_t(tt),V_l_t(tt));
        phi_ll_t(tt)   = ext(V_l_t(tt),V_h_t(tt));
    end
    cond=max(max(abs(hh_in(2:end-1)-phi_hh_t(2:end-1)),abs(ll_in(2:end-1)-phi_ll_t(2:end-1))));
    phi_hh_t(2:end-1)=greed*phi_hh_t(2:end-1)+(1-greed)*hh_in(2:end-1);
    phi_ll_t(2:end-1)=greed*phi_ll_t(2:end-1)+(1-greed)*ll_in(2:end-1);
    cond;
end

%% Main Figures
nameplot='UCI';
calendar=(datetime(2020,1,1)+caldays(1:1000)); 
calendar=datenum(datetime(2020,1,1)+caldays(1:1000)); 

figure(1);
plot(time,[Daleth_se; u_h_se],'LineWidth',3);
legend('Virability','Quarantine Activity Loss'); grid on;  datetick('x','mmm-yy'); axis tight;
orient landscape
set(gca,'LooseInset',get(gca,'TightInset'));
saveas(gcf,[nameplot '_pol'],'pdf');
% f = gcf;
% exportgraphics(f,[nameplot '_pol.pdf']);

% Plots
figure(2);
subplot(2,2,1);
plot(calendar,S_t/mil,'LineWidth',3); grid on; xlabel('Days'); ylabel('Succeptible (millions)'); hold on;
plot(calendar,S_h_t/mil,'LineWidth',3); 
plot(calendar,S_l_t/mil,'LineWidth',3); 
legend('All','high','low'); datetick('x','mmm-yy'); axis tight;  

subplot(2,2,2);
plot(calendar,I_t/mil,'LineWidth',3); grid on; xlabel('Days'); ylabel('Ill (millions)'); datetick('x','mmm-yy'); axis tight;  hold on;
subplot(2,2,3);
plot(calendar,R_t/mil,'LineWidth',3); grid on; xlabel('Days'); ylabel('Recovered (millions)'); datetick('x','mmm-yy'); axis tight;  hold on;
subplot(2,2,4);
plot(calendar,D_t/mil,'LineWidth',3); grid on; xlabel('Days'); ylabel('Dead (millions)'); datetick('x','mmm-yy'); axis tight; hold on;
subplot(2,2,2);
plot(calendar,0*calendar+N*uci_threshold.*UCI_g(time)/mil,'LineWidth',2,'LineStyle',':','Color','k'); 
legend('Ill','UCI beds');
datetick('x','mmm-yy'); axis tight;  
set(gca,'LooseInset',get(gca,'TightInset'));
saveas(gcf,[nameplot '_all'],'pdf');


figure(3)
plot(calendar,delta_t,'LineWidth',3); grid on; xlabel('Days'); ylabel('Deaths rate'); axis tight; hold on;
datetick('x','mmm-yy'); axis tight;  

figure(4)
plot(calendar,0*calendar+max(I_t-N*uci_threshold.*UCI_g(time'),0)/mil,'LineWidth',2); grid on; xlabel('Days'); ylabel('Ill Not Hospitalized (million)'); axis tight; hold on;
datetick('x','mmm-yy'); axis tight;  

figure(5)
plot(calendar,p_t','LineWidth',2); grid on; xlabel('Days'); ylabel('Infection Rate'); axis tight; hold on;
plot(calendar,p_h_t','LineWidth',2,'LineStyle','--'); grid on; xlabel('Days'); ylabel('Infection Rate'); axis tight; hold on;
plot(calendar,p_l_t','LineWidth',2,'LineStyle',':'); 
legend('Average','high','low');
datetick('x','mmm-yy'); axis tight;  

figure(6)
plot(calendar,((Daleth_se.*Daleth_h*m_t.*S_h_t+Daleth_se.*Daleth_l*m_t.*S_l_t)./I_t)'*100,'LineWidth',2); grid on; xlabel('Days'); ylabel('Infection Inflow/Outflow Rate (%)'); axis tight; hold on;
plot(calendar,0*calendar+gamma*100,'LineWidth',2,'LineStyle','--'); grid on; xlabel('Days'); 
legend('Inflow','Outflow')
datetick('x','mmm-yy'); axis tight;  

%% Value Function Block
figure(7)
%plot(calendar,([V_h_t V_l_t V_I V_R*ones(days,1) V_D*ones(days,1)]'),'LineWidth',2); grid on; xlabel('Days'); ylabel('Value'); axis tight; hold on;
plot(calendar,([V_h_t V_l_t V_I V_R*ones(days,1)]'),'LineWidth',2); grid on; xlabel('Days'); ylabel('Value'); axis tight; hold on;
legend('high','low','infected','recovered');
datetick('x','mmm-yy'); axis tight;  

figure(8)
plot(calendar,([V_h_t-V_l_t]'),'LineWidth',2); grid on; xlabel('Days'); ylabel('Value Difference'); axis tight; hold on;
legend('Value High-Low');
datetick('x','mmm-yy'); axis tight;  

figure(9)
plot(calendar,[1-phi_hh_t 1-phi_ll_t]','LineWidth',2); grid on; xlabel('Days'); ylabel('Switching Probability'); axis tight; hold on;
legend('high to low','low to high');
datetick('x','mmm-yy'); axis tight;  

% With dates
disp(['Num Deaths:' num2str(D_t(end))]);

% figure(9)
% plot(calendar,[1-phi_hh_t 1-phi_ll_t]','LineWidth',2); grid on; xlabel('Days'); ylabel('Switching Probability'); axis tight; hold on;
% legend('high to low','low to high');
% datetick('x','mmm-yy'); axis tight;