#include "TTree.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooExponential.h"
#include "RooAddPdf.h"
#include "RooGaussian.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooExtendPdf.h"
#include "RooStats/SPlot.h"
#include "RooStats/ModelConfig.h"
#include "RooGenericPdf.h"
#include "RooFormulaVar.h"
#include "RooBernstein.h"
#include "RooMinuit.h"
#include "RooCurve.h"
using namespace RooFit;
using namespace RooStats;

void PlotTest()
{
 TFile *file = new TFile("results_combined.root");
 RooWorkspace *w = (RooWorkspace *)gDirectory->Get("combined"); // workspace 
  //cout <<"------------------------WorkSpacePrint--------------------------------" <<endl;  
  //w->Print();

  RooAbsData *data = w->data("obsData");    
  //cout <<"---------------------obsData--------------------------------" <<endl;
  //data->Print();

  ModelConfig *sbModel = (ModelConfig *)w->obj("ModelConfig"); // model name 
  cout <<"---------------------ModelConfig--------------------------------" <<endl;

  sbModel->SetName("S+B model");
  sbModel->SetSnapshot(*sbModel->GetParametersOfInterest());
  //cout << "sbModelPrint" << "-----" << endl;
  //sbModel->Print();

  cout << "-------------GetPdf-----------------------"<< endl;
  sbModel->GetPdf()->Print("v");

  ModelConfig *bModel = (ModelConfig *)w->obj("");
  if (!bModel) {
    bModel = (ModelConfig *)sbModel->Clone();
    bModel->SetName("B model");
    RooRealVar *var =
        dynamic_cast<RooRealVar *>(bModel->GetParametersOfInterest()->first());
    if (!var)
      return;
    double oldval = var->getVal();
    var->setVal(0);
    bModel->SetSnapshot(RooArgSet(*var));
    var->setVal(oldval);
  }
   
   RooTrace::active(kTRUE);
   RooTrace::verbose(kTRUE);

  //do fit 
  const RooArgSet *poiSet = sbModel->GetParametersOfInterest();
  RooRealVar *poi = (RooRealVar*)poiSet->first();
  std::cout << "POI initial value:   " << poi->GetName() << " = " << poi->getVal() << std::endl;
  double poihat = 0;
  Info("", " Doing a first background only fit to the observed data --------------");

  RooArgSet constrainParams;
  if (sbModel->GetNuisanceParameters())
      constrainParams.add(*sbModel->GetNuisanceParameters());
  RooStats::RemoveConstantParameters(&constrainParams);
 
  // fit to data
  RooFitResult *fitres = sbModel->GetPdf()->fitTo(*data, InitialHesse(false), Hesse(false),
  Minimizer(ROOT::Math::MinimizerOptions::DefaultMinimizerType().c_str(),"Migrad"),
  Strategy(0), PrintLevel(0), Constrain(constrainParams), Save(true), 
  Offset(RooStats::IsNLLOffset()));

  if (fitres->status() != 0)
    Warning("", " Fit still failed - continue anyway.....");
  cout << "poihat" << endl;
  poihat = poi->getVal();
  std::cout << "Best Fit value : " << poi->GetName() << " = " << poihat
	    << " +/- " << poi->getError() << std::endl;

  cout << "minNLL:" << fitres->minNll() << endl;

  cout << "--------------------fitres:---------------" << endl;
  fitres->Print("V");
  cout << "-------------------NuisanceParameters:---------" << endl;
  sbModel->GetNuisanceParameters()->Print("V"); 
  // save best fit value in the poi snapshot
  sbModel->SetSnapshot(*sbModel->GetParametersOfInterest());
  std::cout << "snapshot of B Model " << sbModel->GetName()
	    << " is set to the best fit value" << std::endl;

   cout << "floatParsFinal" << "-------" << endl;
  fitres->floatParsFinal().Print("s");


  auto pars = fitres->floatParsFinal();
                RooRealVar* p;
                for (int i = 0; i < pars.getSize(); i++) {
                    p = (RooRealVar*)pars.at(i);
                    TString name = p->getTitle().Data();
                   // if (name.Contains("gamma"))
                   //     continue;
                    cout << p->getTitle() << endl;
                    cout << p->getVal() << endl;
                    cout << p->getError() << endl;
                }

   RooCategory sample("channelCat", "channelCat");  // sample
   sample.defineType("set1");               
  

  RooPlot *framex = w->var("obs_x_set1")->frame(Title("Set1"));
  data->plotOn(framex,Cut("channelCat==channelCat::set1"));
  sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),ProjWData(sample,*data),VisualizeError(*fitres));
  //sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),Components("L_x_xe127_set1_overallSyst_x_Exp,sig"),VisualizeError(*fitres),LineColor(kRed));
  //sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),Components("xe127_set1_nominal"),ProjWData(sample,*data),LineColor(kGreen));
  //sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),Components("xe127_set1_epsilon"),ProjWData(sample,*data),LineColor(kPink));
  //sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),Components("set1_model"),ProjWData(sample,*data),LineColor(kYellow));
  //bModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),Components("L_x_xe127_set1_overallSyst_x_Exp,sig"),VisualizeError(*fitres),LineColor(kRed));
  sbModel->GetPdf()->plotOn(framex,Slice(sample,"set1"),ProjWData(sample,*data));
  
  data->plotOn(framex,Cut("channelCat==channelCat::set1"));
  auto canv = new TCanvas("bestfitSet1", "bestfitSet1", 800, 800);
  framex->Draw("same");
  canv->Draw();

  TFile *ff = new TFile("result_test.root","recreate");
  w->Write();
  fitres->Write();
  ff->Close(); 
}