#include <TH1D.h>
#include <TCanvas.h>
#include <TRandom3.h>
#include <TStopwatch.h>
#include <TSystem.h>
#include <TF1.h>

#include <RooRealVar.h>
#include <RooWorkspace.h>
#include <RooDataSet.h>
#include <RooPlot.h>
#include <RooFitResult.h>
#include <RooArgSet.h>
#include <RooAbsReal.h>
#include <RooMinuit.h>
#include <RooSimultaneous.h>
#include <RooCategory.h>

#include <ROOT/RDataFrame.hxx>

#include <RooStats/HistFactory/HistFactorySimultaneous.h>
#include <RooStats/HistFactory/Measurement.h>
#include <RooStats/HistFactory/MakeModelAndMeasurementsFast.h>
#include <RooStats/HistFactory/HistoToWorkspaceFactoryFast.h>

#include <RooStats/ModelConfig.h>
#include <RooStats/ProfileLikelihoodCalculator.h>
#include <iostream>

using namespace std;
using namespace ROOT;
using namespace RooFit;
using namespace RooStats;
using namespace HistFactory;

void createbkgOnlyToyData(Int_t seed);

void doFit(Int_t seed, Double_t *fitVal);

void createSigShape();

void makeFlatHisto();

const Bool_t doStatErr = false; //flag controlling whether Beeston Barlow is used
const Bool_t doPlot = false; //flag controlling is plot graphics are drawn

const Int_t ntoys = 1000;

Int_t fitLow = -5;
Int_t fitHigh = 5;
Int_t nbins = 40;

Int_t genYield = 10000;

void troubleShootHF_2Channel() 
{  
  TStopwatch sw;

  Double_t fitVal = 0;

  TFile *toyResultsFile = new TFile("toyResults_2Channel_10k.root", "RECREATE");
  
  TTree *resTree = new TTree("results","");
  resTree->Branch("fitVal", &fitVal);

  sw.Start();

  createSigShape();
  makeFlatHisto();

  for(Int_t i=1; i<=ntoys; i++)
  {
    if( (i-ntoys)%50 == 0)
    cout<<"####  On toy "<<i<<" of "<<ntoys<<endl;

    createbkgOnlyToyData(i);

    doFit(i, &fitVal);
    
    resTree->Fill();
  }
  cout<<"Done. Time = "<<sw.RealTime()<<endl;
  
  toyResultsFile->Write();
  toyResultsFile->Close();

}

void createSigShape()
{
  TF1 *sig = new TF1("sig", "gaus(0)", fitLow, fitHigh);
  sig->SetParameter(0, 1); //constant
  sig->SetParameter(1, 0); //mean
  sig->SetParameter(2, 1.0); //sigma

  TFile *fileOut = new TFile("sigShape.root", "RECREATE");

  gRandom->SetSeed(1);

  TH1D *sigShape = new TH1D("sigShape","signal shape", nbins, fitLow, fitHigh);

  sigShape->FillRandom("sig", genYield);

  fileOut->Write();
  fileOut->Close();
}

void createbkgOnlyToyData(Int_t seed)
{
  TF1 *bkg = new TF1("bkg", "1.0", -5, 5);

  TH1D *dataHist_Chan1 = new TH1D("data_Chan1", "", nbins, fitLow, fitHigh);
  TH1D *dataHist_Chan2 = new TH1D("data_Chan2", "", nbins, fitLow, fitHigh);

  gSystem->Exec("mkdir -p data_2Channel/");
  TFile *fileOut = new TFile(Form("data_2Channel/data_%d.root", seed), "RECREATE");

  gRandom->SetSeed(seed);

  dataHist_Chan1->FillRandom("bkg", genYield);
  dataHist_Chan2->FillRandom("bkg", genYield);

  dataHist_Chan1->Write();
  dataHist_Chan2->Write();

  fileOut->Close();
  delete dataHist_Chan1;
  delete dataHist_Chan2;
}

void makeFlatHisto()
{
  TFile *fileOut = new TFile("flatShape.root","RECREATE");
  TH1D *flatHist = new TH1D("flat","", nbins, fitLow, fitHigh);

  for(Int_t i=1; i<=nbins; i++)
  {
    flatHist->SetBinContent(i, genYield*1.0/nbins);
  }
  fileOut->Write();
  fileOut->Close();
}
void doFit(Int_t seed, Double_t *fitVal)//, Int_t *status)
{
  const char *histFileName = Form("data_2Channel/data_%d.root", seed);
  
  //*********CONTROL VERBOSITY*****************************************
  RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);

  RooMsgService::instance().getStream(2).minLevel = RooFit::PROGRESS;
  //*******************************************************************

  //***********************SET UP HISTFACTORY**************************
  //Model: Channel 1 has sig+bkg, Channel 2 has only background. Bkg shape is data driven in the fit
  Measurement meas("meas", "test");
  meas.SetOutputFilePrefix("results_2Channel/test");
  meas.SetExportOnly(false);

  meas.SetPOI("nSig");

  meas.SetLumi(1.0);
  meas.SetLumiRelErr(0.00001); 
  meas.AddConstantParam("Lumi");

  Channel chan1("Chan1");
  chan1.SetData("data_Chan1", histFileName);
  chan1.SetStatErrorConfig(0.05, "Poisson"); 

  Channel chan2("Chan2");
  chan2.SetData("data_Chan2", histFileName);
  chan2.SetStatErrorConfig(0.05, "Poisson"); 

  //###################Set up signal sample for Channel 1######################################
  Sample sig_Chan1("sig_Chan1");

  sig_Chan1.SetInputFile("sigShape.root");
  
  sig_Chan1.SetHistoName("sigShape");

  if(doStatErr) sig_Chan1.ActivateStatError();
  sig_Chan1.SetNormalizeByTheory(kFALSE);

  sig_Chan1.AddNormFactor("sigNorm_Chan1", 1.0/genYield, 0., 1.); //normalization of #hist
  meas.AddConstantParam("sigNorm_Chan1");

  sig_Chan1.AddNormFactor("nSig", 10, -genYield, genYield);// the POI

  chan1.AddSample(sig_Chan1);

  //#########################Set up background sample for Channel1#######################################
  Sample bkg_Chan1("bkg");

  bkg_Chan1.SetInputFile("flatShape.root");
  bkg_Chan1.SetHistoName("flat");

  bkg_Chan1.SetNormalizeByTheory(kFALSE);

  bkg_Chan1.AddNormFactor("bkgNorm_Chan1", 1.0/genYield, 0., 1.);
  
  meas.AddConstantParam("bkgNorm_Chan1");

  bkg_Chan1.AddNormFactor("nBkg_Chan1", genYield, -2*genYield, 2*genYield);

  bkg_Chan1.AddShapeFactor("sf");

  chan1.AddSample(bkg_Chan1);

  //#########################Set up background sample for Channel2#######################################
  Sample bkg_Chan2("bkg");

  bkg_Chan2.SetInputFile("flatShape.root");
  bkg_Chan2.SetHistoName("flat");

  bkg_Chan2.SetNormalizeByTheory(kFALSE);

  bkg_Chan2.AddNormFactor("bkgNorm_Chan2", 1.0/genYield, 0., 1.);
  
  meas.AddConstantParam("bkgNorm_Chan2");

  bkg_Chan2.AddNormFactor("nBkg_Chan2", genYield, -2*genYield, 2*genYield);

  bkg_Chan2.AddShapeFactor("sf");

  chan2.AddSample(bkg_Chan2);
  
  //##################################################
  meas.AddChannel(chan1);
  meas.AddChannel(chan2);
  meas.CollectHistograms();
  //###########################Write to a workspac######
  RooWorkspace *w = RooStats::HistFactory::HistoToWorkspaceFactoryFast::MakeCombinedModel(meas);
  
  gSystem->Exec("mkdir -p fitResults_2Channel/");

  TString wsFileName = Form("fitResults_2Channel/ws_%d.root", seed);
  w->writeToFile(wsFileName.Data(), true);

  //#####Read workspace and do fit
  TFile *ws_file = TFile::Open(wsFileName.Data(), "READ");
  RooWorkspace* ws = (RooWorkspace*) ws_file->Get("combined");
  RooStats::ModelConfig* myModelConfig = (RooStats::ModelConfig*) ws->obj("ModelConfig");

  RooSimultaneous *simPdf = (RooSimultaneous *) myModelConfig->GetPdf(); 

  RooAbsData* myData = ws->data("obsData");
  ws_file->Close();

  // Access the confidence interval on the parameter of interest (POI).
  RooRealVar* myPOI = (RooRealVar*) myModelConfig->GetParametersOfInterest()->first();
  
  RooAbsReal *myNLL1 = simPdf->createNLL(*myData, Offset(kTRUE));
  RooMinuit *minuit = new RooMinuit(*myNLL1);

  minuit->setPrintLevel(-1);
  minuit->setNoWarn();

  minuit->setStrategy(2);
  minuit->fit("smh");

  *fitVal =  myPOI->getVal();

  double fitErr = myPOI->getError();
  
  double fitPull = (*fitVal - 0)/fitErr;

  if(doPlot)
  {
    RooRealVar *obs_Chan1 = (RooRealVar *) myModelConfig->GetObservables()->find("obs_x_Chan1");
    RooPlot *histFrame_Chan1 = obs_Chan1->frame(Bins(nbins));

    RooRealVar *obs_Chan2 = (RooRealVar *) myModelConfig->GetObservables()->find("obs_x_Chan2");
    RooPlot *histFrame_Chan2 = obs_Chan2->frame(Bins(nbins));

    RooCategory *cat = (RooCategory*)(myModelConfig->GetObservables())->find("channelCat");

    myData->plotOn(histFrame_Chan1, DataError(RooAbsData::Poisson), Cut("channelCat==channelCat::Chan1"), DrawOption("ZP"), Name("data_Chan1") );
    myData->plotOn(histFrame_Chan2, DataError(RooAbsData::Poisson), Cut("channelCat==channelCat::Chan2"), DrawOption("ZP"), Name("data_Chan2") );

    simPdf->plotOn(histFrame_Chan1, Slice(*cat, "Chan1"), ProjWData(*cat, *myData), LineWidth(2), LineColor(4), Name("fit_Chan1"));
    simPdf->plotOn(histFrame_Chan2, Slice(*cat, "Chan2"), ProjWData(*cat, *myData), LineWidth(2), LineColor(4), Name("fit_Chan2"));

    if(*fitVal > 0)
      simPdf->plotOn(histFrame_Chan1, Slice(*cat, "Chan1"), ProjWData(*cat, *myData), Components("L_x_sig_Chan1_Chan1_overallSyst_x_Exp"), LineWidth(2), LineColor(2), Range(int(fitLow), int(fitHigh)));

    simPdf->plotOn(histFrame_Chan1, Slice(*cat, "Chan1"), ProjWData(*cat, *myData), Components("L_x_bkg_Chan1_overallSyst_x_Exp_x_Chan1_sf_shapeFactor"), LineWidth(2), LineColor(6), Range(int(fitLow), int(fitHigh)));
    simPdf->plotOn(histFrame_Chan2, Slice(*cat, "Chan2"), ProjWData(*cat, *myData), Components("L_x_bkg_Chan2_overallSyst_x_Exp_x_Chan2_sf_shapeFactor"), LineWidth(2), LineColor(6), Range(int(fitLow), int(fitHigh)));

    new TCanvas();
    histFrame_Chan1->Draw();

    new TCanvas();
    histFrame_Chan2->Draw();

  }

  //gSystem->Exec(Form("rm -f data/data_%d.root", seed));
  gSystem->Exec(Form("rm -f %s", wsFileName.Data()));
}