using nlohmann::json;


// HistFactory toy script to perform 1D and 2D fits with/witout BB

//................

// Activate most recent ROOT release 6.28/06 on lxplus :

      // source /cvmfs/sft.cern.ch/lcg/views/setupViews.sh LCG_104a x86_64-centos7-gcc11-opt

// Activate fixed ROOT version 6.29/01 with nightlies on lxplus :

      // source /cvmfs/sft.cern.ch/lcg/views/dev3/Tue/x86_64-centos7-gcc12-opt/setup.sh

//................

// Arguments of the main function 'toyModel_1D_and_2D_wBB' :

    // 1) bool makeToyHistos       ==> set 'true' to make initial template histograms of data, signal and background    
    // 2) bool makeReNormHistos    ==> set 'true' to renormalise fitting templates to yield obtained without activating Beeston Barlow 
    // 3) bool BB                  ==> set 'true' to activate Beeston Barlow treatment of template statistical uncertainties
    // 4) int DIM                  ==> set 2 or 1 for 1D or 2D template fit


//................

// Run toymodel in four steps with ROOT 6.28/06: 

  // (1) make initial fitting templates: 

        // root 'toyModel_1D_and_2D_wBB.C(true,false,false,2)'

  // (2) run template fit without BB to get yields to re-normalise templates      

        // root 'toyModel_1D_and_2D_wBB.C(false,false,false,2)'

  // (3) make re-normalised templates    

        // root 'toyModel_1D_and_2D_wBB.C(false,true,false,2)'

  // (4) run template fit WITH BB to get final result     

        // root 'toyModel_1D_and_2D_wBB.C(false,false,true,2)'

//................

// Run toymodel in four steps with ROOT 6.29/01 with nightlies: 

  // (1) make initial fitting templates: 

        // root --web=off 'toyModel_1D_and_2D_wBB.C(true,false,false,2)'

  // (2) run template fit without BB to get yields to re-normalise templates      

        // root --web=off 'toyModel_1D_and_2D_wBB.C(false,false,false,2)'

  // (3) make re-normalised templates    

        // root --web=off 'toyModel_1D_and_2D_wBB.C(false,true,false,2)'

  // (4) run template fit WITH BB to get final result     

        // root --web=off 'toyModel_1D_and_2D_wBB.C(false,false,true,2)'





//============================================================================
// Create 1D histogram templates
//============================================================================


void createToyHistos_1D(){


  cout << "\n>>>>>> Create histogram templates..\n" << endl;


  gRandom->SetSeed(1356);

  // Save histograms in this root file
  TFile* histoFile = TFile::Open("histoFile.root", "RECREATE");

  // Set number of events in signal and background (N_Sig = N_BG = N_DATA/2)
  int N_evt = 10000; 

  // Set min/max range of observable x
  double x_min = 3000;
  double x_max = 6000;

  // Number of bins
  int N_bins = 15;

  // Functions used to fill histograms
  TF1* func1 = new TF1("func1","1 + 0.01 * x",x_min, x_max); // background
  TF1* func2 = new TF1("func2","0.2 * exp(-( (x-4000.)*(x-4000.)/(2.0*200.*200.) ) )", x_min, x_max); // signal
  
  // Histograms
  TH1D* histoBG = new TH1D("histoBG","background histo", N_bins, x_min, x_max);
  TH1D* histoSig = new TH1D("histoSig","signal histo", N_bins, x_min, x_max);
  TH1D* histoData = new TH1D("histoData","data histo", N_bins, x_min, x_max);
  
  // define observable x 
  double x_bg; // background
  double x_sig; // signal

  // Fill background and signal histograms
  for( int i = 0; i <  N_evt ; ++i){
    x_bg = func1->GetRandom();
    x_sig = func2->GetRandom();
    histoBG->Fill(x_bg);
    histoSig->Fill(x_sig);
  }
  
  // Fill data histogram
  for( int i = 0; i <  N_evt ; ++i){
    x_bg = func1->GetRandom();
    histoData->Fill( x_bg );
    x_sig = func2->GetRandom();
    histoData->Fill( x_sig );
  }

  cout << "Initial histogram yields : " << endl;
  cout << "N_Data = " << histoData->Integral() << endl;
  cout << "N_Sig = " << histoSig->Integral() << endl;
  cout << "N_BG = " << histoBG->Integral() << "\n" << endl;


  histoFile->Write();
  histoFile->Close();

}

//============================================================================
// Create 2D histogram templates
//============================================================================


void createToyHistos_2D(){


  cout << "\n>>>>>> Create histogram templates..\n" << endl;

  gRandom->SetSeed(1356);

  // save templates in this root file
  TFile* histoFile = TFile::Open("histoFile.root", "RECREATE");

  // Set number of events in signal and background (N_Sig = N_BG = N_DATA/2)
  int N_evt = 10000; 

  // Set min/max range of observable x
  double xy_min = 3000;
  double xy_max = 6000;

  // Number of bins
  int N_bins = 15;

  // Functions used to fill histograms
  TF2* func1 = new TF2("func1","1 + 0.01 * x * 0.01 * y*y",xy_min, xy_max, xy_min, xy_max); // background
  TF2* func2 = new TF2("func2","0.2 * exp(-( (x-4000.)*(x-4000.)/(2.0*200.*200.) + (y-5000.)*(y-5000.)/(2.0*300*300)) )", 3000, 6000, 3000, 6000); // signal
  
  // Histograms
  TH2D* histoBG = new TH2D("histoBG","background histo", N_bins, xy_min, xy_max, N_bins, xy_min, xy_max);
  TH2D* histoSig = new TH2D("histoSig","signal histo", N_bins, xy_min, xy_max, N_bins, xy_min, xy_max);
  TH2D* histoData = new TH2D("histoData","data histo", N_bins, xy_min, xy_max, N_bins, xy_min, xy_max);
  
  // define observables x and y 
  double x_bg, y_bg;
  double x_sig, y_sig;

  // Fill background and signal histograms
  for( int i = 0; i < N_evt; ++i){
    func1->GetRandom2(x_bg,y_bg);
    func2->GetRandom2(x_sig,y_sig);
    histoBG->Fill(x_bg,y_bg);
    histoSig->Fill(x_sig,y_sig);
  }

  // define observables x and y for data
  double xd_bg, yd_bg;
  double xd_sig, yd_sig;
  
  // Fill data histogram
  for( int i = 0; i < N_evt; ++i){
    func1->GetRandom2(xd_bg, yd_bg);
    histoData->Fill( xd_bg, yd_bg );

    func2->GetRandom2(xd_sig, yd_sig);
    histoData->Fill( xd_sig, yd_sig );
  }

  cout << "Initial histogram yields : " << endl;
  cout << "N_Data = " << histoData->Integral() << endl;
  cout << "N_Sig = " << histoSig->Integral() << endl;
  cout << "N_BG = " << histoBG->Integral() << "\n" << endl;


  histoFile->Write();
  histoFile->Close();

}


//============================================================================
// Create re-normalised templates
//============================================================================


void createReNormHistos(const std::string histoFileName){

  cout << "\n>>>>>> Re-normalise histogram templates..\n" << endl;

  std::ifstream ifs("output_noBB.json");
  json jf = json::parse(ifs);

  Double_t scale_sig = jf["histoSig"];
  Double_t scale_BG = jf["histoBG"];

  TFile* histoFile = TFile::Open(histoFileName.c_str(), "READ");

  TH1D* histoSig = (TH1D*)histoFile->Get("histoSig")->Clone();
  TH1D* histoBG = (TH1D*)histoFile->Get("histoBG")->Clone();


  TFile* histoFile_reNorm = TFile::Open("histoFile_reNorm.root", "RECREATE");

  histoSig->Sumw2();
  histoSig->Scale(scale_sig/histoSig->Integral());
  
  histoBG->Sumw2();
  histoBG->Scale(scale_BG/histoBG->Integral());


  cout << "Renormalised histogram yields : " << endl;
  cout << "N_Sig = " << histoSig->Integral() << endl;
  cout << "N_BG = " << histoBG->Integral() << "\n" << endl;
 
  histoSig->SetDirectory(histoFile_reNorm );
  histoBG->SetDirectory(histoFile_reNorm );

  histoFile_reNorm->Write();
  histoFile_reNorm->Close();

}


//============================================================================
// Load templates
//============================================================================


std::vector<TH1D*> getHistos(const std::string histoFileName, bool BB){

  cout << "\n>>>>>> Loading histograms from file..\n" << endl;

  TFile* histoFile = TFile::Open(histoFileName.c_str(), "READ");
  std::vector<TH1D*> histos;
  

  TH1D* histoSig = (TH1D*)histoFile->Get("histoSig")->Clone();
  histos.push_back( histoSig );


  TH1D* histoBG = (TH1D*)histoFile->Get("histoBG")->Clone();
  cout << histoBG->GetName() << endl;
  histos.push_back( histoBG );

  return histos;
}


//============================================================================
// Perform toymodel fit
//============================================================================


void toyModel_1D_and_2D_wBB(bool makeToyHistos = false, bool makeReNormHistos = false, bool BB = false, int DIM = 2){
  
  cout << "\n>>>>>> HistFactory toy fit in " << DIM << "D..\n" << endl;

  TString fitType;

  if(BB){
    fitType = "yesBB";
    std::cout << "Beeston-Barlow is activated " << std::endl;
  }

  if(!BB){ 
    fitType = "noBB";
    std::cout << "Beeston-Barlow is NOT activated" << std::endl;
  }


  // --------------------------------------
  // Make histograms 
  // --------------------------------------

  // make initial templates
  if( makeToyHistos && DIM == 2){
    TCanvas* c = new TCanvas();
    createToyHistos_2D();
    return;
  }
  if( makeToyHistos && DIM == 1){
    TCanvas* c = new TCanvas();
    createToyHistos_1D();
    return;
  }


  // --------------------------------------
  // Load histograms 
  // --------------------------------------

  // file with initial templates
  std::string histoFileName = "histoFile.root";

  // make re-normalised templates based on fit result witout BB
  if( makeReNormHistos ){
    TCanvas* c = new TCanvas();
    createReNormHistos(histoFileName);
    return;
  }

  // use re-normlised templates
  if(BB) histoFileName = "histoFile_reNorm.root";

  // load templates 
  std::vector<TH1D*> templateHistos = getHistos(histoFileName, BB);


  // --------------------------------------
  // Setup measurement & channel 
  // --------------------------------------

  // Intital yileds
  double N_sig_init = 10000.0;
  double N_bg_init  = 10000.0;

  //bool constraining = true;
  bool constraining = false;

  //--- Constraining background to signal
  double frac_bg_vs_sig = 0.0;
  double lim_bg_vs_sig_min = 0.0;
  double lim_bg_vs_sig_max = 0.0;


  // -- Define the measurement
  std::cout << "==> Defining the measurement" << std::endl;

  RooStats::HistFactory::Measurement B2D0MuNuMeas("B2D0MuNu", "B2D0MuNu fit"); 
  
  B2D0MuNuMeas.SetExportOnly(true); // Tells histfactory to not run the fit
  B2D0MuNuMeas.SetPOI("num_histoSig"); // set to Bogus parma. of interest 
  B2D0MuNuMeas.SetLumi(1.0);
  B2D0MuNuMeas.SetLumiRelErr(0.1);
  
  // -- Define the channel (in our case, there is only one)
  RooStats::HistFactory::Channel B2D0MuNu("BMCorr");
  B2D0MuNu.SetStatErrorConfig(1e-5,"Poisson"); //Configure the use of bin-by-bin statistical uncertainties
  B2D0MuNu.SetData( "histoData", "histoFile.root" );

  // -- Add the samples to the channel (= the individual templates)  
  for( TH1D* histo : templateHistos ){
    
    RooStats::HistFactory::Sample sample(histo->GetName(),histo->GetName(),histoFileName);
    
    if(BB) sample.ActivateStatError();
    sample.SetNormalizeByTheory(false);
    sample.AddNormFactor( "norm_"+(std::string)histo->GetName(), 1/histo->Integral(), 1/histo->Integral(), 1/histo->Integral() ); // fix normalisation
    sample.AddNormFactor( "num_"+(std::string)histo->GetName(), histo->Integral(),  1.0,  histo->Integral()*2 ); // free yield
    
    // add histograms to channel
    B2D0MuNu.AddSample( sample );
  }

  // add channel to measurements
  B2D0MuNuMeas.AddChannel( B2D0MuNu );
  B2D0MuNuMeas.CollectHistograms();


  // --------------------------------------
  // Setting up workspace
  // --------------------------------------

  std::cout << " " << std::endl;
  std::cout << "\n==> Setting up workspace" << "\n" << std::endl;

  // -- Define workspace
  RooWorkspace* workSpace;

  // --------- Define MakeModelAndMeasurementFast
  workSpace = RooStats::HistFactory::MakeModelAndMeasurementFast(B2D0MuNuMeas, {.binnedFitOptimization=false});

  // Get model manually
  RooStats::ModelConfig* modelConf = (RooStats::ModelConfig*)workSpace->obj("ModelConfig");
  RooSimultaneous* model = (RooSimultaneous*)modelConf->GetPdf();

  ((RooRealVar*)(modelConf->GetNuisanceParameters()->find("Lumi")))->setConstant(true); //fix the Lumi to not have it also fitted
  ((RooRealVar*)(modelConf->GetNuisanceParameters()->find("norm_histoSig")))->setConstant(true); //fix norm
  ((RooRealVar*)(modelConf->GetNuisanceParameters()->find("norm_histoBG")))->setConstant(true); //fix norm

  RooArgSet* obs = (RooArgSet*)modelConf->GetObservables();
  RooCategory *idx = (RooCategory*)obs->find("channelCat");
  RooAbsData* data = (RooAbsData*)workSpace->data("obsData");
  
  RooRealVar* obs_x = (RooRealVar*)obs->find("obs_x_BMCorr");
  RooRealVar* obs_y = (RooRealVar*)obs->find("obs_y_BMCorr");


  // make BB work with newer ROOT versions ('HistFactorySimultaneous' removed - instead use directly 'RooSimultaneous')
  RooSimultaneous* totalPdf = new RooSimultaneous( *model );

  // make nll model to data
  RooAbsReal* nll =  totalPdf->createNLL( *data , RooFit::Offset(kTRUE) );

  // BB in likelihood
  RooStats::HistFactory::RooBarlowBeestonLL* bbnll = new RooStats::HistFactory::RooBarlowBeestonLL("bbnll", "bbnll", *nll); 
  bbnll->setPdf( totalPdf );
  bbnll->setDataset( data );
  bbnll->initializeBarlowCache(); 


  // --------------------------------------
  // Perform fit
  // --------------------------------------

  
  // --------- Minimization code
  RooMinimizer* minuitHF = new RooMinimizer( *bbnll ) ; // nll
  minuitHF->setErrorLevel(0.5);
  minuitHF->optimizeConst(1);
  minuitHF->migrad() ; // Minimization
  minuitHF->hesse() ; // Hessian error analysis  

  RooFitResult *fitResult = minuitHF->save("Result","Result"); 
  fitResult->Print("v");


  // --------------------------------------
  // Extract fit result
  // --------------------------------------

  // Signal fit result
  RooRealVar* fitResult_sig = (RooRealVar*)fitResult->floatParsFinal().find("num_histoSig");
  Double_t N_sig_fit = fitResult_sig->getVal();
  Double_t err_N_sig_fit = fitResult_sig->getError();

  // BG fit result
  RooRealVar* fitResult_bg = (RooRealVar*)fitResult->floatParsFinal().find("num_histoBG");
  Double_t N_BG_fit = fitResult_bg->getVal();
  Double_t err_yield_bg = fitResult_bg->getError();
  
  
  std::cout << " " << std::endl;
  std::cout << "Fit result (BB= : " << BB << ")" << std::endl;
  std::cout << " " << std::endl;
  std::cout << "N_sig_fit = " << N_sig_fit << " +- " << err_N_sig_fit << std::endl;
  std::cout << "N_BG_fit = " << N_BG_fit << " +- " << err_yield_bg << std::endl;
  std::cout << " " << std::endl;


  // --------------------------------------
  // Saving yields to json file
  // --------------------------------------

  const double nSig = N_sig_fit;
  const double nBG = N_BG_fit;

  std::ofstream output_file("output_" + fitType + ".json");

  json outJson;

  outJson["histoSig"] = nSig;
  outJson["histoBG"] = nBG;

  output_file << outJson;
  output_file.close(); 



  // --------------------------------------
  // Plotting fit result
  // --------------------------------------


  std::cout << " " << std::endl;
  std::cout << "==> Start plotting.. " << std::endl;
  std::cout << " " << std::endl;

  double min = 3000; 
  double max = 6000; 


  std::vector<std::string> varNames; 

  if (DIM == 1) varNames.push_back("obs_x_BMCorr");

  if (DIM == 2) {
    varNames.push_back("obs_x_BMCorr");
    varNames.push_back("obs_y_BMCorr");
  }

  RooPlot* frame;
  RooPlot* frame_tot;
  RooPlot* pull;
  TLegend* legend;

   RooPlot *ffr;

  for (int i = 0; i < DIM; i++){

    if (varNames[i]=="obs_x_BMCorr") frame = obs_x->frame();
    if (varNames[i]=="obs_y_BMCorr") frame = obs_y->frame();

  
    data->plotOn(frame, RooFit::DataError(RooAbsData::Poisson), RooFit::MarkerSize(0.4), RooFit::DrawOption("ZP"), RooFit::Name("pdf_Data"));
    model->plotOn(frame, RooFit::ProjWData(*idx, *data), RooFit::Name("pdf_model"), RooFit::Components("*histoSig*,*histoBG*"), RooFit::DrawOption("L"), RooFit::LineColor(kGreen)); 
    model->plotOn(frame, RooFit::ProjWData(*idx, *data), RooFit::Name("pdf_sig"), RooFit::Components("*histoSig*"), RooFit::DrawOption("L"), RooFit::LineColor(kBlue)); 
    model->plotOn(frame, RooFit::ProjWData(*idx, *data), RooFit::Name("pdf_bg"), RooFit::Components("*histoBG*"), RooFit::DrawOption("L"), RooFit::LineColor(kRed)); 

    frame->GetYaxis()->SetRangeUser(0.0, 5.5e3);
    frame->GetXaxis()->SetTitle(" "); 
    frame->SetTitle(" ");

    if (varNames[i]=="obs_x_BMCorr") legend = new TLegend(0.70, 0.4, 0.90, 0.96);
    if (varNames[i]=="obs_y_BMCorr") legend = new TLegend(0.15, 0.4, 0.35, 0.96);
    legend->SetBorderSize(0);
    legend->SetFillColorAlpha(kWhite, 0.0);
    legend->AddEntry(frame->findObject("pdf_Data"), "Data", "lep");
    legend->AddEntry(frame->findObject("pdf_model"), "Fit", "l");
    legend->AddEntry(frame->findObject("pdf_sig"), "Signal", "l");
    legend->AddEntry(frame->findObject("pdf_bg"), "Background", "l");
    legend->Draw();

   // - - - - - - combined plot

    if (varNames[i]=="obs_x_BMCorr") frame_tot = obs_x->frame();
    if (varNames[i]=="obs_y_BMCorr") frame_tot = obs_y->frame();
    data->plotOn(frame_tot, RooFit::DataError(RooAbsData::Poisson), RooFit::MarkerSize(0.4), RooFit::DrawOption("ZP"), RooFit::Name("pdf_Data"));
    model->plotOn(frame_tot, RooFit::ProjWData(*idx, *data), RooFit::Name("total"), RooFit::Components("*histoSig*,*histoBG*"), RooFit::DrawOption("L"), RooFit::LineColor(kGreen));
    

    // - - - - - - pull plot

    if (varNames[i]=="obs_x_BMCorr") pull = obs_x->frame();
    if (varNames[i]=="obs_y_BMCorr") pull = obs_y->frame();
    RooHist* h1 = frame_tot->pullHist();
    pull->addPlotable(h1,"E1");

    pull->GetXaxis()->SetLabelSize(0.13);
    pull->GetYaxis()->SetLabelSize(0.13);
    pull->GetYaxis()->SetNdivisions(5);
    pull->SetTitle(" ");
    pull->GetXaxis()->SetTitle(" "); 
    pull->SetMaximum(4.5);
    pull->SetMinimum(-4.5);
    TLine l_up(min, 2., max, 2.);
    TLine l_m(min, 0., max, 0.);
    TLine l_down(min, -2., max, -2.);
    l_up.SetLineColor(kRed);
    l_down.SetLineColor(kRed);
    l_m.SetLineStyle(3);


    // - - - - - -  plot fit and pull plots together

    TCanvas c("Fit_and_pull", "Fit and pull", 900, 600);
    c.Divide(1,2);

    c.cd(1);
    gPad->SetPad(0.0, 0.20, 1.0, 1.0);
    gPad->SetRightMargin(0.025);
    gPad->SetTopMargin(0.02);
    frame->Draw();
    legend->Draw("same");

    c.cd(2);
    gPad->SetPad(0.0, 0.00, 1.0, 0.25);
    if (varNames[i]=="obs_x_BMCorr") pull->GetXaxis()->SetTitle("obs_x_BMCorr"); 
    if (varNames[i]=="obs_y_BMCorr") pull->GetXaxis()->SetTitle("obs_y_BMCorr"); 
    pull->GetXaxis()->SetTitleOffset(0.9);
    pull->GetXaxis()->SetTitleSize(0.17);
    
    gPad->SetRightMargin(0.025);
    gPad->SetBottomMargin(0.40);
    pull->Draw("e1");
    l_up.Draw("same");
    l_down.Draw("same");
    l_m.Draw("same");

    c.Update();

    TString fitName = "Fit_";
    if (varNames[i]=="obs_x_BMCorr") fitName += "mcorr";
    if (varNames[i]=="obs_y_BMCorr") fitName += "mpipi";
    fitName += ".pdf";

    c.SaveAs( fitName ); 

}


}