%% This script replicates subsample panel data estimation results in:
% "Social Distancing, Vaccination and Evolution of Covid-19 Transmission Rates in Europe"
% by A. Chudik, M. H. Pesaran and A. Rebucci, January 2022. 
%
%
% Data is loaded first, either using a previously saved matlab file dataload1.mat or using the following input files : 
% - Covid-19 data in confirmed_cases_txt,
% - WHO Covid-19 data for Spain and France from OxCGRT_latest_21Dec2021.csv,  
% - Population data in population.xlsx, 
% - Oxford indices data in OxCGRT_latest_21Dec2021.csv downloaded from https://github.com/OxCGRT/covid-policy-tracker/tree/master/data
% - vaccination data in owid-covid-data_21Dec2021.xlsx downloaded from https://github.com/owid/covid-19-data/tree/master/public/data
% - variants data in covid-variants.csv downloaded from https://github.com/owid/covid-19-data/tree/master/public/data/variants (on 4 Jan 2022)
%
% See also ReadMe.txt for further information.

clc 
clear
load_mat_data=true;  % if set to true, data is loaded from matdata.mat
                      % if set to false, data is loaded from individual files outlined above (this takes longer)  

if load_mat_data
    load matdata; 
else

% load data

%% covid-19 data   
    disp('loading covid-19 data from local txt file...')
 
 Ac = readcell('confirmed_cases_txt', 'DatetimeType','text');
 % population data
 disp('loading population data from xlsx file...')
 Ap = readcell('population.xlsx');
  % delete missing population entries
  j=2;
  for i=2:size(Ap,1)
      if ismissing(Ap{j,5})
          Ap(j,:)=[];
      else
          j=j+1;
      end
  end
Ap=reorderAP(Ap,Ac); % to make sure order is identical;

 Ac=correct_french_data(Ac); % if executed, it will replace JHU french data with WHO data - please check and update (if needed) the WHO data range in this file!
 Ac=correct_spanish_data(Ac); % if executed, it will replace JHU french data with WHO data - please check and update (if needed) the WHO data range in this file!
   
% add unreported country aggregates to Ac
Ac=add_aggregates(Ac);
Ap=add_aggregates(Ap);

 
%% Oxford data
   disp('loading Oxford indices data from csv file...');
  Ago = readcell('OxCGRT_latest_21Dec2021.csv');
  % clear regions data to save space
  ind=false(size(Ago,1),1);
  for i=2:size(Ago,1)
      a=Ago{i,3};
      if size(a,2)>1
        ind(i,1)=true;%not(ismissing(a));
      end
  end
   Ago(ind,:)=[];
  % replace names to match JUH data labels
  Ago=replace_nameso(Ago,'South Korea','Korea, South');
  Ago=replace_nameso(Ago,'Myanmar','Burma');
  Ago=replace_nameso(Ago,'Democratic Republic of Congo','Congo (Kinshasa)');
  Ago=replace_nameso(Ago,'Congo','Congo (Brazzaville)');
  Ago=replace_nameso(Ago,'Czech Republic','Czechia');
  Ago=replace_nameso(Ago,'Slovak Republic','Slovakia');
  Ago=replace_nameso(Ago,'Taiwan','Taiwan*');
  Ago=replace_nameso(Ago,'United States','US');

dt=Ac(1,5:end); % date
cl=country_list(Ac); % country list

% further processing of Oxford data
disp('processing Oxford data');

indX=[51,59]; % selector of variables - check the file to make sure correct columns are selected!
Xnames={['Stringency Index']; ['Ec. Support Index']};

Ao=cell(size(Ac));
Ao(2:end,1)=cl;
Ao(1,5:end)=dt;
for i=1:size(cl,1)

   country_to_find=cl{i,1};
   [sp,foundc]=find_go_start(Ago,country_to_find);
   if foundc
    for j=1:size(dt,2)
        date_to_find=dt{1,j};
        [ox,~,foundd,sp1]= find_go_sp(Ago,country_to_find,date_to_find, sp, indX);
        for h=1:size(ox,2)
            Ao(i+1,j+4,h)=ox(1,h);
        end
        if foundd
            sp=sp1;
        end
    end
   else
   % disp(['...country ',country_to_find,' not found in the Oxford database']);
   end
end
clear Ago;

%% Owid vaccination data
  disp('loading owid-covid-data from local xlsx file...')
 Awo = readcell('owid-covid-data_21Dec2021.xlsx');
 Awo(:,1:2)=[]; % first two columns are not relevant;

disp('processing OWID data'); 
   % replace names to match JUH data labels
  Awo=replace_nameso(Awo,'South Korea','Korea, South');
  Awo=replace_nameso(Awo,'Myanmar','Burma');
  Awo=replace_nameso(Awo,'Democratic Republic of Congo','Congo (Kinshasa)');
  Awo=replace_nameso(Awo,'Congo','Congo (Brazzaville)');
  Awo=replace_nameso(Awo,'Taiwan','Taiwan*');
  Awo=replace_nameso(Awo,'United States','US');
 
indX=[43-2]; % selector of variables - check the file to make sure correct columns are selected!
Xnames(end+1,1)={['Fully Vaccinated per 100']};

Aw=cell(size(Ac));
Aw(2:end,1)=cl;
Aw(1,5:end)=dt;

datecolumn=zeros(size(Awo,1),1);
parfor i=2:size(Awo,1)
    %i
    datecolumn(i,1)=str2double(datestr(Awo{i,2}, 'yyyymmdd'));
end

for i=1:size(cl,1)
   country_to_find=cl{i,1};
 %  i
   [sp,foundc]=find_go_start(Awo,country_to_find);
   if foundc
    for j=1:size(dt,2)
        date_to_find=dt{1,j};
        [wx,~,foundd,sp1]= find_go_spw(Awo,country_to_find,date_to_find, sp, indX, datecolumn);
        for h=1:size(wx,2)
            Aw(i+1,j+4,h)=wx(1,h);
        end
        if foundd
            sp=sp1;
        end
    end
   else
   % disp(['...country ',country_to_find,' not found in the OWID database']);
   end
end

% backward fill missing vaccination variable observations with zeros
for i=2:size(Aw,1)
    if size(Aw{i,5},1)<1
        Aw{i,5}=0; % initial values are set to zero
    end
    for t=6:size(Aw,2)
        if size(Aw{i,t},1)<1
            Aw{i,t}=Aw{i,t-1};
        end
    end
end
clear Awo;

%% variants data
  disp('loading covid variants data from local csv file...')
 Avo = readcell('covid-variants.csv','DatetimeType','text');
 Avo(:,[4,6])=[]; % these columns are not relevant;

disp('processing covid variants data'); 
   % replace names to match JUH data labels
  Avo=replace_nameso(Avo,'South Korea','Korea, South');
  Avo=replace_nameso(Avo,'Myanmar','Burma');
  Avo=replace_nameso(Avo,'Democratic Republic of Congo','Congo (Kinshasa)');
  Avo=replace_nameso(Avo,'Congo','Congo (Brazzaville)');
  Avo=replace_nameso(Avo,'Taiwan','Taiwan*');
  Avo=replace_nameso(Avo,'United States','US');
 
varstr={['Alpha'],['Delta'],['Omicron']};
Xnames(end+1,1)={['Alpha share per 100']};
Xnames(end+1,1)={['Delta share per 100']};
Xnames(end+1,1)={['Omicron share per 100']};
indX=4;

Av=cell(size(Ac));
Av(2:end,1)=cl;
Av(1,5:end)=dt;

datecolumn=zeros(size(Avo,1),1);
parfor i=2:size(Avo,1)
    %i
    datecolumn(i,1)=str2double(datestr(Avo{i,2}, 'yyyymmdd'));
end

for i=1:size(cl,1)
   country_to_find=cl{i,1};
   %i
   [avos,foundc, datecs]=find_sel(Avo,country_to_find,1, datecolumn);
   if foundc
    for vi=1:size(varstr,2)
        [avosv,foundv, datecsv]=find_sel(avos,varstr{vi},3,datecs);
        for j=1:size(dt,2)
            date_to_find=dt{1,j};
            [wx,~,foundd,sp1]= find_go_spw(avosv,country_to_find,date_to_find, 1, indX, datecsv);
            if foundd
                Av(i+1,j+4,vi)=wx;
                sp=sp1;
            end
        end
    end% vi
   else
   % disp(['...country ',country_to_find,' not found in the OWID database']);
   end
end

% backward fill missing vaccination variable observations with zeros
for v=1:size(varstr,2)
    for i=2:size(Av,1)
        if size(Av{i,5,v},1)<1
            Av{i,5,v}=0; % initial values are set to zero
        end
        for t=6:size(Av,2)
            if size(Av{i,t,v},1)<1
                Av{i,t,v}=Av{i,t-1,v};
            end
        end
    end
end
clear Avo avos avosv;
%% add aggregates to all
Ao=add_aggregatesWo(Ao, Ac, Ap);
Aw=add_aggregatesWo(Aw, Ac, Ap);
Av=add_aggregatesWo(Av, Ac, Ap);
Ac=add_aggregatesW(Ac, Ac);
Ap=add_aggregatesW(Ap, Ac);

save matdata;
end %load_mat_data

%% computations
disp(['computing results ...'])

% data transformations
 % get number of countries/regions and time dimensions
n=size(Ac,1)-1; %-1 due to the first row being descriptions
T=size(Ac,2)-4; %-4 due to descriptions

cl=country_list(Ac);
pl=country_list(Ap);

C=cell2mat(Ac(2:end,5:end));
P=fill_entries(cell2mat(replace_br_with_nans(Ap(2:end,5:end))),pl,cl);

%now compute synthetic R series, assuming gamma
gamma=1/14; % do not change! There are several other instances in the codes that set gamma=1/14;
Rem2=createRe(C,gamma);

% compute number of active infections
I2=C-Rem2;
PT=P*ones(1,size(C,2));
s=(PT-C)./PT;

%compute c and i
c=C./PT;
i2=I2./PT;

% compute daily cases per 100k, 7dma
[dcpm, dCm]=compute_dc(C,PT);

matransform=true;W=14; % Set Moving Average transformation, and estimation window size
sv=50; % minimum threshold for computing cases
MF=2; [b_1wm_mf2, r_1wm_mf2]=compute_trans_rates_hc(MF*c, MF*i2,W,dCm, matransform, sv); 
MF=3; [b_1wm_mf3, r_1wm_mf3]=compute_trans_rates_hc(MF*c, MF*i2,W,dCm, matransform, sv); 
MF=2.5; [b_1wm_mf2_5, r_1wm_mf2_5]=compute_trans_rates_hc(MF*c, MF*i2,W,dCm, matransform, sv); 
MF=5; [b_1wm_mf5, r_1wm_mf5]=compute_trans_rates_hc(MF*c, MF*i2,W,dCm, matransform, sv);  
MF=7; [b_1wm_mf7, r_1wm_mf7]=compute_trans_rates_hc(MF*c, MF*i2,W,dCm, matransform, sv); 

cl=country_list(Ac); % country list
%% panel estimations

k=size(Ao,3);
for i=1:k
    Apom=Ao(2:end,5:end,i);
    Apom=replace_br_with_nans(Apom);
    X(:,:,i)=cell2mat(Apom);
end
    Awpom=Aw(2:end,5:end,1);
    Awpom=replace_br_with_nans(Awpom);
    X(:,:,k+1)=cell2mat(Awpom);
kv=size(Av,3);
for i=1:kv
    Avom=Av(2:end,5:end,i);
    Avom=replace_br_with_nans(Avom);
    X(:,:,k+1+i)=cell2mat(Avom);
end
    
k=k+kv+1;

MF=5; % Multiple factor
sv=100; % for selection of panel estimation sample

% define samples of countries
selected_countries={'Belgium';'France'; 'Germany';'Italy';'Netherlands';'Poland';'Portugal';...
    'Spain';'United Kingdom'};
selected_countries_short={'Bel';'Fra';'Ger';'Ita';'Neth';'Pol';'Por';...
    'Spain';'UK'};


Thvar.data=dcpm; 
Thvar.name=['Confirmed cases per day per 100k population (7day ma)'];
% define grid
gr(1:10,1)=[1:10]/1000;
gr(11:29,1)=0.01+[1:19]/100;
gr(30:37,1)=0.2+[1:8]/10;
gr(38:55,1)=1+[1:18]/2;
gr(56:80,1)=10+[1:25]*2;
Thvar.grid=gr;
Thvar.usegrid=true;

SC.list=selected_countries;
SC.short_names=selected_countries_short;
Data.c=c;
Data.i2=i2;
Data.I2=I2;
Data.Ac=Ac;
Data.b=b_1wm_mf5;
Data.R=r_1wm_mf5;
Data.dates=dt;

% select regressors
Xn(:,:,1)=X(:,:,1); Xnms(1,1)={['Stringency Index']};
Xn(:,:,2)=X(:,:,2); Xnms(2,1)={['Ec. Support']};
Data.X=Xn;
Data.Xnames=Xnms;


do_thresholds=true;

% select time period
Data.c(:,377:end)=NaN;% data until end Jan 2021
save_file='Subsample_panel_estimation_results.xlsx';
lags=[10,12]; % lag orders
panel_estimations_mf(SC,Data, Thvar, sv, save_file,1, do_thresholds, 2,3, lags); 
save_file='Subsample_panel_estimation_results_robustness.xlsx';
lags=[7,14]; % lag orders
panel_estimations_mf(SC,Data, Thvar, sv, save_file,1, do_thresholds, 2,3, lags); 



disp('All done. Please check ReadMe.txt regarding the description of the results xlsx file.')