using namespace RooFit;

void FinalMassFitToysSimultaneousSimple(){

  TFile *myfile2016= new TFile("ToyDataset_2016.root","READ");
  TFile *myfile2017= new TFile("ToyDataset_2017.root","READ");
  TFile *myfile2018= new TFile("ToyDataset_2018.root","READ");
  TTree *mytree2016= (TTree*)myfile2016->Get("modelData");
  TTree *mytree2017= (TTree*)myfile2017->Get("modelData");
  TTree *mytree2018= (TTree*)myfile2018->Get("modelData");

  RooRealVar x("m","m_{inv} (MeV/c^{2})",3400,3800);
  RooDataSet data16("2016","m_{inv} (MeV/c^{2})",x,Import(*mytree2016));
  RooDataSet data17("2017","m_{inv} (MeV/c^{2})",x,Import(*mytree2017));
  RooDataSet data18("2018","m_{inv} (MeV/c^{2})",x,Import(*mytree2018));

  // Define category to distinguish samples
  RooCategory sample("sample","sample");
  sample.defineType("2016");
  sample.defineType("2017");
  sample.defineType("2018");

  // Construct combined dataset in (x,sample)
  RooDataSet data("data","combined data",x,Index(sample),Import("2016",data16),Import("2017",data17),Import("2018",data18));

  RooRealVar mu("mu","mu",3400,3800);
  mu.setVal(3576.);
  mu.setConstant(kTRUE);
  RooRealVar sig("sig","sig",6.5,6.0,10.0);
  sig.setVal(9.);
  sig.setConstant(kTRUE);
  RooRealVar lambda("lambda", "lambda",  0.001, -0.05, 0.05);
  //lambda.setVal(0.001019);
  //lambda.setConstant(kTRUE);

  RooPlot* xframe = x.frame(Title(""));
  data.plotOn(xframe,Name("data"),Binning(40));

  // --- Build DSCB PDF ---
  RooGaussian signal16("signal16", "gausian", x, mu, sig);
  RooGaussian signal17("signal17", "gausian", x, mu, sig);
  RooGaussian signal18("signal18", "gausian", x, mu, sig);

  // background PDFs
  // ------- expon ------- //
  RooExponential bkg16("bkg16", "exponential PDF", x, lambda);
  RooExponential bkg17("bkg17", "exponential PDF", x, lambda);
  RooExponential bkg18("bkg18", "exponential PDF", x, lambda);

  RooRealVar sigFrac16("sigFrac16", "Signal Fraction", 0., 0, 1);
  RooRealVar sigFrac17("sigFrac17", "Signal Fraction", 0., 0, 1);
  RooRealVar sigFrac18("sigFrac18", "Signal Fraction", 0., 0, 1);
  RooAddPdf model16("model16", "Signal plus background PDF",  signal16, bkg16, sigFrac16);
  RooAddPdf model17("model17", "Signal plus background PDF",  signal16, bkg17, sigFrac17);
  RooAddPdf model18("model18", "Signal plus background PDF",  signal16, bkg18, sigFrac18);

  // Construct a simultaneous pdf using category sample as index
  RooSimultaneous simModel("simModel","simultaneous pdf",sample);
  RooSimultaneous simBkg("simBkg","simultaneous pdf",sample);

  // Associate model with the physics state and model_ctl with the control state
  simModel.addPdf(model16,"2016");
  simModel.addPdf(model17,"2017");
  simModel.addPdf(model18,"2018");
  simBkg.addPdf(bkg16,"2016");
  simBkg.addPdf(bkg17,"2017");
  simBkg.addPdf(bkg18,"2018");

  RooFitResult* resultBkg = simBkg.fitTo(data,Save());
  RooFitResult* resultSigBkg = simModel.fitTo(data,Save());

  TCanvas* c1 = new TCanvas("c1","c1");
  simModel.plotOn(xframe,Name("model"));
  simModel.plotOn(xframe,Slice(sample,"2016"),Slice(sample,"2017"),Slice(sample,"2018"));
  simModel.plotOn(xframe, ProjWData(data), LineColor(kRed));
  xframe->SetTitle("");
  xframe->Draw();

  RooPlot* frame16 = x.frame(Bins(40),Title("2016"));
  data.plotOn(frame16,Cut("sample==sample::2016"));
  simModel.plotOn(frame16,Slice(sample,"2016"),Components("bkg16"),ProjWData(sample,data),LineStyle(kDashed),LineColor(kGreen));
  simModel.plotOn(frame16,Slice(sample,"2016"),ProjWData(sample,data));
  RooPlot* frame17 = x.frame(Bins(40),Title("2017"));
  data.plotOn(frame17,Cut("sample==sample::2017"));
  simModel.plotOn(frame17,Slice(sample,"2017"),Components("bkg17"),ProjWData(sample,data),LineStyle(kDashed),LineColor(kGreen));
  simModel.plotOn(frame17,Slice(sample,"2017"),ProjWData(sample,data));
  RooPlot* frame18 = x.frame(Bins(40),Title("2018"));
  data.plotOn(frame18,Cut("sample==sample::2018"));
  simModel.plotOn(frame18,Slice(sample,"2018"),Components("bkg18"),ProjWData(sample,data),LineStyle(kDashed),LineColor(kGreen));
  simModel.plotOn(frame18,Slice(sample,"2018"),ProjWData(sample,data));
  TCanvas* c = new TCanvas("c","c",1500,400);
  c->Divide(3);
  c->cd(1); gPad->SetLeftMargin(0.15); frame16->GetYaxis()->SetTitleOffset(1.4); frame16->Draw();
  c->cd(2); gPad->SetLeftMargin(0.15); frame17->GetYaxis()->SetTitleOffset(1.4); frame17->Draw();
  c->cd(3); gPad->SetLeftMargin(0.15); frame18->GetYaxis()->SetTitleOffset(1.4); frame18->Draw();

  Double_t qValue;
  Double_t pValue;
  qValue = 2*(resultBkg->minNll() - resultSigBkg->minNll()) > 0.0001 ? 2*(resultBkg->minNll() - resultSigBkg->minNll()) : 0;
  pValue = 0.5*(1. - TMath::Erf(sqrt(qValue/2.)));

  RooFitResult* resultSigBkg16 = model16.fitTo(data16,Save());
  RooFitResult* resultSigBkg17 = model17.fitTo(data17,Save());
  RooFitResult* resultSigBkg18 = model18.fitTo(data18,Save());
  RooFitResult* resultBkg16 = bkg16.fitTo(data16,Save());
  RooFitResult* resultBkg17 = bkg17.fitTo(data17,Save());
  RooFitResult* resultBkg18 = bkg18.fitTo(data18,Save());

  printf("lambda is %f ± %f\n",lambda.getVal(),lambda.getError());
  printf("mean is %f\n",mu.getVal());
  printf("p-value is %.12f\n",pValue);
  printf("significance is %f\n",RooStats::PValueToSignificance(pValue));
  printf("----------------------------------------\n");
  printf("2016 bkg model minNll is %f\n",resultBkg16->minNll());
  printf("2016 sig+bkg model minNll is %f\n", resultSigBkg16->minNll());
  printf("2017 bkg model minNll is %f\n",resultBkg17->minNll());
  printf("2017 sig+bkg model minNll is %f\n", resultSigBkg17->minNll());
  printf("2018 bkg model minNll is %f\n",resultBkg18->minNll());
  printf("2018 sig+bkg model minNll is %f\n", resultSigBkg18->minNll());
  printf("sum of 16/17/18 bkg model minNll is %f\n",resultBkg16->minNll()+resultBkg17->minNll()+resultBkg18->minNll());
  printf("sum of 16/17/18 sig+bkg model minNll is %f\n",resultSigBkg16->minNll()+resultSigBkg17->minNll()+resultSigBkg18->minNll());
  printf("simultaneous fit bkg model minNll is %f\n",resultBkg->minNll());
  printf("simultaneous fit sig model minNll is %f\n", resultSigBkg->minNll());
}