#include <cassert>
#include <vector>
#include "Math/Interpolator.h"
#include "TMultiGraph.h"
#include "TVirtualX.h"
#include "TCanvas.h"
#include "TGraph.h"
#include "TGraphErrors.h"
#include "TLegend.h"
#include "TError.h"
#include "TColor.h"
#include "TMath.h"
#include "TStyle.h"
#include "TH1F.h"

#include<stdio.h>
#include<iostream>

using namespace std;
const int N = 200;

void bfunc ( double* par, double* ttemp, double* stemp, double* itemp, double* rtemp) {

   // fitting the SIR Model to the data

   double p1 = par[0];
   double p2 = par[0] / par[1];
   double p3 = par[2];
   double p4 = par[3];
   double dtime  = 1.0;

   stemp[0] = 1.0;
   itemp[0] = p3;
   rtemp[0] = 0.0;
   ttemp[0] = 0.0;

   for (int j=1; j<N; j++) {

      stemp[j] = stemp[j-1] - p1*stemp[j-1]* itemp[j-1]*dtime;
      itemp[j] = itemp[j-1] + (p1*stemp[j-1]*itemp[j-1] - p2*itemp[j-1])*dtime;
      rtemp[j] = rtemp[j-1] + p2*itemp[j-1]*dtime;
      ttemp[j] = j*dtime;
   }

   // re-scale the distributions
   for (int j=0; j<N; j++) {
     stemp[j] *= p4;
     itemp[j] *= p4;
     rtemp[j] *= p4;

   }

}

double func_act(double* x, double* par) {

   // fitting the SIR Model to the data
   // active

   double stemp[N];
   double itemp[N];
   double rtemp[N];
   double ttemp[N];
   
   bfunc (par, ttemp, stemp, itemp, rtemp);

   /*
   ROOT::Math::Interpolator inter(N, ROOT::Math::Interpolation::kCSPLINE);
   inter.SetData(N, ttemp, itemp);
   double result = inter.Eval ( *x );
   */

   TGraph *gtemp = new TGraph (N, ttemp, itemp);
   double result = gtemp->Eval ( *x );
   delete gtemp;

   return result;
}


double func_rec(double* x, double* par) {

   // fitting the SIR Model to the data
   // recovered

   double stemp[N];
   double itemp[N];
   double rtemp[N];
   double ttemp[N];

   bfunc (par, ttemp, stemp, itemp, rtemp);

   TGraph *gtemp = new TGraph (N, ttemp, rtemp);
   double result = gtemp->Eval ( *x );
   delete gtemp;

   return result;
}

void Covid19_Togo () {

//TCanvas *c1 = new TCanvas("c1","SIR Model",0,0,200,1);

TCanvas *c1 = new TCanvas("c1","Covid-19 Togo",200,10,700,500);

const double population = 8.0e+6;
const int Ndata = 35;


/* Daily data */
/*
double data_infected  [Ndata] = {1, 8, 3, 4, 2, 2, 3, 1, 1, 3, 2, 4, 0, 2, 3, 1, 1, 3, 14 };
double data_time      [Ndata] = {0, 13, 14, 15, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31};
double data_recovered [Ndata] = {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 6, 3, 0, 0, 7, 0, 1, 2, 3};
double data_dead      [Ndata] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0};
double data_imported  [Ndata] = {1, 3, 0, 0, 0, 2, 3, 0, 0, 2, 0, 3, 0, 0, 2, 2, 1, 0, 2};
*/

/* Integrated data up to that day */
double data_infected  [Ndata] = {1, 9, 12, 16, 18, 20, 23, 24, 25, 28, 30, 34, 34, 36, 39, 40, 41, 44, 58, 65, 70, 73, 76, 76, 76, 77, 77, 81, 81, 83, 84, 84, 84, 86, 88 };
double data_time      [Ndata] = {0, 13, 14, 15, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47};
double data_recovered [Ndata] = {0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 10, 10, 10, 17, 17, 18, 20, 23, 23, 23, 24, 25, 25, 29, 29, 32, 35, 45, 48, 49, 52, 53, 56, 56};
double data_dead      [Ndata] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 5, 5, 5, 5, 6, 6, 6};
double data_active    [Ndata] = {1, 8, 11, 15, 17, 19, 22, 23, 23, 26, 22, 23, 23, 24, 20, 20, 20, 20, 32, 39, 44, 46, 48, 48, 44, 45, 42, 43, 31, 30, 30, 27, 25, 24, 26};


const int Nlabel = 10;
const char *data_date[Nlabel]  = {"3/6/2020","3/16/2020","3/26/2020","4/5/2020",
      "4/15/2020","4/25/2020","5/5/2020", "5/15/2020","5/25/2020", "6/5/2020" };

double XactErr[Ndata];
double YactErr[Ndata];
double YrecErr[Ndata];

for (int j=0; j<Ndata; j++) {

 YactErr        [j] = sqrt ( data_active[j] ) / population;
 YrecErr        [j] = sqrt( data_recovered [j] ) / population;

 data_infected  [j] /= population;
 data_recovered [j] /= population;
 data_dead      [j] /= population;
 data_active    [j] /= population;

 XactErr[j] = 0.25;
}

TGraph *dinf, *drec, *ddea;
TGraphErrors *dact;
dinf  = new TGraph(Ndata,data_time,data_infected);
//dinf->SetMinimum(-0.5);  
dinf->GetYaxis()->SetTitle("Fractions of people");
dinf->SetLineWidth(2);
dinf->SetLineColor(2);
for (int j=1;j<=Nlabel;j++) dinf->GetXaxis()->ChangeLabel(j,30,-1,-1,-1,-1,data_date[j-1]);
dinf->SetMarkerColor(2);
dinf->SetMarkerStyle(20);
dinf->SetTitle("Covid-19 Togo Data");
dinf->Draw("AP");

drec = new TGraphErrors(Ndata,data_time,data_recovered,XactErr,YrecErr);
drec->SetLineWidth(2);
drec->SetLineColor(3);
drec->SetMarkerStyle(20);
drec->SetMarkerColor(3);
drec->Draw("PSame");

ddea = new TGraph(Ndata,data_time,data_dead);
ddea->SetLineWidth(2);
ddea->SetLineColor(6);
ddea->SetMarkerStyle(20);
ddea->SetMarkerColor(6);
ddea->Draw("PSame");

dact = new TGraphErrors(Ndata,data_time,data_active,XactErr,YactErr);
dact->SetLineWidth(2);
dact->SetLineColor(4);
dact->SetMarkerStyle(20);
dact->SetMarkerColor(4);
dact->Draw("PSame");

TF1 *actfit = new TF1("actfit",func_act,0.,47.,4);
double par[4];

par[0] = 0.24; // 0.3;    //30;
par[1] = 2.13; //1.0;     //1.0;
par[2] = 4.37e-08; 1.91e-7; //1.25e-7;
par[3] = 1.0;  //1.0;

actfit->SetParameters(par[0], par[1], par[2], par[3]);
actfit->SetParLimits(0, 0.15, 70.0);
actfit->SetParLimits(1, 0.5, 20.0);
actfit->SetParLimits(2, 0.075e-7, 1.25e-5);
actfit->FixParameter(3, 1.0);


TVirtualFitter::SetMaxIterations(80000);
TVirtualFitter::SetDefaultFitter("Minuit");
ROOT::Math::MinimizerOptions::SetDefaultStrategy(2);
ROOT::Math::MinimizerOptions::SetDefaultTolerance(0.0002);

actfit->SetLineColor(4);
actfit->SetLineStyle(4);
actfit->SetLineWidth(4);
TFitResultPtr acFit = dact->Fit("actfit", "SMEXO");
actfit->Draw("Same");

par[0] = acFit->Value(0);
par[1] = acFit->Value(1);
par[2] = acFit->Value(2);
par[3] = acFit->Value(3);

TF1 *recfit = new TF1("recfit",func_rec,0.,47.,4);
recfit->SetParameters(par[0], par[1], par[2], par[3]);

recfit->SetParLimits(0, 0.0, 70.0);
recfit->SetParLimits(1, 0.5, 10.0);
recfit->SetParLimits(2, 0.075e-7, 1.25e-3);
recfit->SetParLimits(3, 0, 1.1);

recfit->SetLineColor(3);
recfit->SetLineStyle(4);
recfit->SetLineWidth(4);
TFitResultPtr reFit = drec->Fit("recfit", "SMEXO");
recfit->Draw("Same");

TF1 *closure = new TF1("closure",func_act,0.,47.,4);
par[0] = reFit->Value(0);
par[1] = reFit->Value(1);
par[2] = reFit->Value(2);
par[3] = reFit->Value(3);
closure->SetParameters(par[0], par[1], par[2], par[3]);
closure->SetLineColor(7);
closure->SetLineStyle(3);
closure->SetLineWidth(4);
closure->Draw("Same");

auto legend = new TLegend(0.2,0.6,0.40,0.85);
//legend->SetFillStyle(1001);

legend->SetHeader("Covid-19 Togo","C"); // option "C" allows to center the header
TLegendEntry * entry = legend->AddEntry(dinf,"Infected","p");
//entry->SetMarkerStyle(20);
//entry->SetMarkerSize(40);
legend->AddEntry(dact,"Active","p");
legend->AddEntry(drec,"Recovered","p");
legend->AddEntry(ddea,"Dead","p");
legend->AddEntry(actfit,"SIR Model Fit to Active","l");
legend->AddEntry(recfit,"SIR Model Fit to Recovered","l");
legend->AddEntry(closure,"Closure Test","l");
legend->Draw();

//c1->Update();

return;
}