#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "RooPlot.h"
#include "TText.h"
#include "TArrow.h"
#include "TFile.h"
#include "RooDataHist.h"

using namespace RooFit;

void toyMC(){
      gROOT->Reset();
      gSystem->Load("libRooFit");

      Float_t 
            range0=4, 
            range1=11,
            x0=4.1, 
            x1=10.8,
            max_nsig=0;
      Int_t 
            nbins = 30, 
            nMCsamples = 5e1,
            nMCevents = 101,
            iMCtestSample[2] = {0, rand()%nMCsamples},
            count = 0;

      vector<int> fails;

      RooRealVar mass("mass","Generated mass variable",range0,range1);
      mass.setBins(nbins); 

      RooRealVar a0_gen("a0_gen","Generated a0", 3.2901e-01);
      RooRealVar a1_gen("a1_gen","Generated a1",-9.0915e-01);
      RooRealVar a2_gen("a2_gen","Generated a2",-6.7814e-02);
      RooRealVar a3_gen("a3_gen","Generated a3", 3.2901e-01);
      RooRealVar a4_gen("a4_gen","Generated a4", 5.3439e-02);
      RooRealVar a5_gen("a5_gen","Generated a5", 9.9368e-03);
      RooChebychev cpol_gen("cpol_gen","Generated Chebyshev polynomial", mass, RooArgList(a0_gen, a1_gen, a2_gen, a3_gen, a4_gen, a5_gen));
      
      RooRealVar mean_fit("mean_fit","Fitted BW mean", 4.75, 4.5, 5.0);
      RooRealVar sigma_fit("sigma_fit","Fitted BW width", 0.25, 0.001, 0.75);
      RooBreitWigner bw_fit("bw_fit", "Fitted gauss", mass, mean_fit, sigma_fit);

      RooRealVar a0_fit("a0_fit","Fitted a0", 3.2901e-01,-1.,1.);
      RooRealVar a1_fit("a1_fit","Fitted a1",-9.0915e-01,-1.,1.);
      RooRealVar a2_fit("a2_fit","Fitted a2",-6.7814e-02,-1.,1.);
      RooRealVar a3_fit("a3_fit","Fitted a3", 3.2901e-01,-1.,1.);
      RooRealVar a4_fit("a4_fit","Fitted a4", 5.3439e-02,-1.,1.);
      RooRealVar a5_fit("a5_fit","Fitted a5", 9.9368e-03,-1.,1.);
      RooChebychev cpol_fit("cpol_fit","Fitted Chebyshev polynomial", mass, RooArgList(a0_fit, a1_fit, a2_fit, a3_fit, a4_fit, a5_fit));

      RooRealVar nsig("nsig","N sgn evts", 0, 0, 10);
      RooRealVar nbkg("nbkg","N bkg evts", 101, 10, 200);

      RooAddPdf fit("fit", "bw+cpol", RooArgList(bw_fit, cpol_fit), RooArgList(nsig, nbkg));

      RooMCStudy *mcstudy = new RooMCStudy(cpol_gen, mass, FitModel(fit), Silence(), Binned(),
                                                           FitOptions(PrintLevel(-1),
                                                           Minimizer("Minuit2","minimize"),
                                                           Range(x0, x1),
                                                           Save()
                                                           // EvalErrorWall(kFALSE)
                                                           ));

      mcstudy->generateAndFit(nMCsamples, nMCevents, kTRUE); //last argument orders to save generated datasets

      // Iterating over all MC samples and corresponding fits to pick up the test sample and count the number of breaches
      for(int j = 0; j<nMCsamples; j++){
            RooAbsData *gendata_j = mcstudy->genData(j); 
            const RooFitResult *fitres_j = mcstudy->fitResult(j);
            RooArgSet const& pars_j = fitres_j->floatParsFinal();
            auto& nsig_j = static_cast<RooRealVar&>(pars_j["nsig"]);
            if(nsig_j.getVal()>3)
                  count++, fails.push_back(j);
            if(nsig_j.getVal()>max_nsig)
                  max_nsig = nsig_j.getVal();
            cout << nsig_j.getVal() << "+-" << nsig_j.getError() << endl;
      }

      // Let's pick up a random exotic MC sample fit 
      iMCtestSample[0]=fails[rand()%fails.size()];
      RooAbsData *gendata_i = mcstudy->genData(iMCtestSample[0]); 
      const RooFitResult *fitres_i = mcstudy->fitResult(iMCtestSample[0]);
      RooArgSet const& pars_i = fitres_i->floatParsFinal();

      // Record fit values for a random exotic MC sample fit 
      auto& mean_i  = static_cast<RooRealVar&>(pars_i["mean_fit"]);
      auto& sigma_i = static_cast<RooRealVar&>(pars_i["sigma_fit"]);
      auto& a0_i    = static_cast<RooRealVar&>(pars_i["a0_fit"]);
      auto& a1_i    = static_cast<RooRealVar&>(pars_i["a1_fit"]);
      auto& a2_i    = static_cast<RooRealVar&>(pars_i["a2_fit"]);
      auto& a3_i    = static_cast<RooRealVar&>(pars_i["a3_fit"]);
      auto& a4_i    = static_cast<RooRealVar&>(pars_i["a4_fit"]);
      auto& a5_i    = static_cast<RooRealVar&>(pars_i["a5_fit"]);
      auto& nsig_i  = static_cast<RooRealVar&>(pars_i["nsig"]);
      auto& nbkg_i  = static_cast<RooRealVar&>(pars_i["nbkg"]);

      // Let's pick up a random test sample
      RooAbsData *gendata_k = mcstudy->genData(iMCtestSample[1]); 
      const RooFitResult *fitres_k = mcstudy->fitResult(iMCtestSample[1]); // get fit result of a toyMC no. i
      RooArgSet const& pars_k = fitres_k->floatParsFinal();
      
      // Record fit values for the respective fit 
      auto& mean_k  = static_cast<RooRealVar&>(pars_k["mean_fit"]);
      auto& sigma_k = static_cast<RooRealVar&>(pars_k["sigma_fit"]);
      auto& a0_k    = static_cast<RooRealVar&>(pars_k["a0_fit"]);
      auto& a1_k    = static_cast<RooRealVar&>(pars_k["a1_fit"]);
      auto& a2_k    = static_cast<RooRealVar&>(pars_k["a2_fit"]);
      auto& a3_k    = static_cast<RooRealVar&>(pars_k["a3_fit"]);
      auto& a4_k    = static_cast<RooRealVar&>(pars_k["a4_fit"]);
      auto& a5_k    = static_cast<RooRealVar&>(pars_k["a5_fit"]);
      auto& nsig_k  = static_cast<RooRealVar&>(pars_k["nsig"]);
      auto& nbkg_k  = static_cast<RooRealVar&>(pars_k["nbkg"]);
     
      RooPlot *frame1 = mcstudy->plotParam(nsig, Range(0.,max_nsig+1), Bins(1));
      RooPlot *frame2 = mcstudy->plotParam(nbkg, Bins(nbins));

      RooPlot *frame3 = mass.frame(Bins(nbins));
      frame3->SetTitle(Form("Test toy sample No. %i (threshold exceed)",iMCtestSample[0]));
      gendata_i->plotOn(frame3);
      mean_fit.setVal(mean_i.getVal());
      sigma_fit.setVal(sigma_i.getVal());
      a0_fit.setVal(a0_i.getVal());
      a1_fit.setVal(a1_i.getVal());
      a2_fit.setVal(a2_i.getVal());
      a3_fit.setVal(a3_i.getVal());
      a4_fit.setVal(a4_i.getVal());
      a5_fit.setVal(a5_i.getVal());
      nsig.setVal(nsig_i.getVal());
      nbkg.setVal(nbkg_i.getVal());
      fit.plotOn(frame3);
      fit.plotOn(frame3,Components(RooArgSet(bw_fit)),LineColor(kYellow-2),LineStyle(kDashed));
      fit.plotOn(frame3,Components(RooArgSet(cpol_fit)),LineColor(kRed),LineStyle(kDashed));
      
      RooPlot *frame4 = mass.frame(Bins(nbins));
      frame4->SetTitle(Form("Test toy sample No. %i (random)",iMCtestSample[1]));
      gendata_k->plotOn(frame4);
      mean_fit.setVal(mean_k.getVal());
      sigma_fit.setVal(sigma_k.getVal());
      a0_fit.setVal(a0_k.getVal());
      a1_fit.setVal(a1_k.getVal());
      a2_fit.setVal(a2_k.getVal());
      a3_fit.setVal(a3_k.getVal());
      a4_fit.setVal(a4_k.getVal());
      a5_fit.setVal(a5_k.getVal());
      nsig.setVal(nsig_k.getVal());
      nbkg.setVal(nbkg_k.getVal());
      fit.plotOn(frame4);
      fit.plotOn(frame4,Components(RooArgSet(bw_fit)),LineColor(kYellow-2),LineStyle(kDashed));
      fit.plotOn(frame4,Components(RooArgSet(cpol_fit)),LineColor(kRed),LineStyle(kDashed));

      TCanvas *c = new TCanvas("mcstudy", "mcstudy", 800, 800);
      c->Divide(2, 2);
      c->cd(1);
      frame1->Draw();
      c->cd(2);
      frame2->Draw();
      c->cd(3);
      frame3->Draw();
      c->cd(4);
      frame4->Draw();

      cout << "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^" << '\n';
      cout << Form("Fit No.%i paramerters:",iMCtestSample[0]) << endl;
      fitres_i->Print("v");
      
      // for(int i = 0; i < 10; i++)
      //       fitres_i->floatParsFinal().at(i)->Print();

      cout << "Nsig_max:    " << max_nsig << endl;
      cout << "Number of MC samples with a combinatoric peak detected:    " << count << endl;
      cout << "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^" << '\n';
}