%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%   Optimal Treatment Strategy
%
%   This code computes the optimal treatment strategy for an SIS type
%   model. Note, this code is using the dimensionless version of the sytem.
%
%   Inputs:
%   x0=ratio of number of infections to total population
%   R0=standard reproduction number=N*beta/alpha
%   A=ratio of death rate/recovery rate
%   T=maximum time horizon
%   n=number of temporal points
%   w1=weight of infection cost
%   w2=weight of treatment cost
%
%   Ouputs:
%   1. Plot of infections assuming no control strategy and plot of
%   infections with the control strategy.
%   2. Plot of the control strategy.
%   3. Plot of integrand of the Lagrangian assuming no control and plot of
%   the integrand of the Lagrangian with the optimal control.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Inputs
x0=.1;
R0=1.5;
A=.1;
T=20;
n=1000;
w1=1;
w2=1;

%% Construction of temporal mesh
t=linspace(0,T,n);

%% Solution to ODEs without the optimal control
[~,x] = ode45(@(t,x) SIS(t,x,R0,A),t,x0);

%% Solution with suboptimal control
[~,xSub] = ode45(@(t,x) SISSub(t,x,R0,A),t,x0);


%% Solution to ODEs with the optimal control
[~,xControlled] = ode45(@(t,x) SISControlled(t,x,R0,A,w1,w2,T),t,x0);

%% Plotting Infected Status

figure('units','normalized','outerposition',[0 0 1 1]);

plot(t,x,t,xSub,t,xControlled,'linewidth',2)
xlabel('$$t$$','interpreter','latex')
ylabel('Infected Population','interpreter','latex')
legend('No Control','Constant Control','Optimal Control','interpreter','latex','fontsize',40)
ax=gca;
ax.FontSize = 40;
set(gcf,'color','w');

%% Plotting Control

figure('units','normalized','outerposition',[0 0 1 1]);

plot(t,zeros(1,n),t,ones(1,n),t,w1/(2*w2)*(T-t).*xControlled','linewidth',2)
xlabel('$$t$$','interpreter','latex')
ylabel('Control Parameter','interpreter','latex')
legend('No Control','Constant Control','Optimal Control','interpreter','latex','fontsize',40)
ax=gca;
ax.FontSize = 40;
set(gcf,'color','w');

%% Plotting Functional

figure('units','normalized','outerposition',[0 0 1 1]);

plot(t,w1*x,t,w1*xSub+w2*ones(n,1).^2,t,w1*xControlled+w2*(w1/(2*w2)*(T-t').*xControlled).^2,'linewidth',2)
xlabel('$$t$$','interpreter','latex')
ylabel('Cost Density','interpreter','latex')
legend('No Control','Constant Control','Optimal Control','interpreter','latex','fontsize',40)
ax=gca;
ax.FontSize = 40;
set(gcf,'color','w');

%% SIS ODE System
function xdot = SIS(t,x,R0,A)
xdot=R0*x*(1-x)-x-A*x;
end

%% SIS ODE Suboptimal System
function xdot = SISSub(t,x,R0,A)
xdot=R0*x*(1-x)-x-A*x-x;
end

%% SIS Controlled ODE System
function xdot = SISControlled(t,x,R0,A,w1,w2,T)
xdot=R0*x*(1-x)-x-A*x+w1*(t-T)/(2*w2)*x^2;
end


