#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooAddPdf.h"

#include <iostream>

using namespace std;
using namespace RooFit;



void TestFitYieldsSimultaneously(Bool_t checkGeneration = false, Int_t genSigW = 6000, Int_t genCombW = 200, Int_t genSigNW = 54000, Int_t genCombNW = 800)
{
  Int_t cpus = 1;
  
  ////////////////////////
  // Model for Wide sample
  ////////////////////////
  // Mass
  Double_t       m_low       = 1790.;
  Double_t       m_high      = 1940.;
  RooRealVar     *m          = new RooRealVar(    "D0_M",       "Invariant mass",             m_low,  m_high);
  RooRealVar     *mean       = new RooRealVar(    "mean",       "Mean of gaussians",          1866.,  m_low,  m_high);
  // Signal
  RooRealVar     *sigma1     = new RooRealVar(    "sigma1",     "sigma1",                     4.5,     2.,     10.);
  RooGaussian    *gausCore   = new RooGaussian(   "gausCore",   "Core gaussian",              *m,     *mean, *sigma1);
  // Combinatorial
  RooRealVar     *slopeComb  = new RooRealVar(    "slopeComb",  "Combinatorial slope",         -0.012, -10.,     -0.0001);
  RooExponential *pdfComb    = new RooExponential("pdfComb",    "Combinatorial pdf",           *m,    *slopeComb);
  // Yields
  RooRealVar     *signalW    = new RooRealVar(    "signalW",    "Signal events Wide",           genSigW,  1., RooNumber::infinity());
  RooRealVar     *combW      = new RooRealVar(    "combW",      "Combinatorial events Wide",    genCombW, 1., RooNumber::infinity());
  // Model
  RooAddPdf      *massModelW = new RooAddPdf(     "massModelW", "g1 + expo",                  RooArgList(*gausCore, *pdfComb), RooArgList(*signalW, *combW));
  


  //////////////////////////
  // Model for NoWide sample
  //////////////////////////
  // Yields
  RooRealVar     *signalNW    = new RooRealVar(    "signalNW",    "Signal events NoWide",         genSigNW,  1., RooNumber::infinity());
  RooRealVar     *combNW      = new RooRealVar(    "combNW",      "Combinatorial events NoWide",  genCombNW, 1., RooNumber::infinity());
  // Model
  RooAddPdf      *massModelNW = new RooAddPdf(     "massModelNW", "g1 + expo",                  RooArgList(*gausCore, *pdfComb), RooArgList(*signalNW, *combNW));



  ////////////////
  // Generate data
  ////////////////
  RooArgSet  *vars          = new RooArgSet("Variables");
  vars->add(*m);
  
  // W
  RooDataSet *thisDatasetW   = massModelW ->generate(RooArgSet(*m), genSigW  + genCombW);
  
  // NW
  m->setMin(1825.);
  m->setMax(1910.);
  RooDataSet *thisDatasetNW  = massModelNW->generate(RooArgSet(*m), genSigNW + genCombNW);
  

  
  //////////////
  // Set binning
  //////////////
  // Double_t binsPerMeV = 1.;
  // m->setBins(binsPerMeV * (Int_t)(m_high - m_low + 0.00000001));
  // m->Print("");



  ////////////////
  // Define ranges
  ////////////////
  m->setRange("FitRange_Wide",   m_low, m_high);
  m->setRange("FitRange_NoWide", 1825., 1910.);

  

  /////////////////////////////////////////
  // Create index category and join samples
  /////////////////////////////////////////

  // Define category to distinguish the two samples events
  RooCategory sample("sample", "sample");
  sample.defineType("Wide");
  sample.defineType("NoWide");

  // Construct combined dataset in (m, sample)
  RooDataHist *binnedW  = thisDatasetW ->binnedClone();
  RooDataHist *binnedNW = thisDatasetNW->binnedClone();
  // RooDataSet  combData("combData", "combined data", *m, Index(sample), Import("Wide", *thisDatasetW), Import("NoWide", *thisDatasetNW));
  RooDataHist combData("combData", "combined data", *m, Index(sample), Import("Wide", *binnedW),      Import("NoWide", *binnedNW));



  //////////////////////////////////////////////
  // Construct a simultaneous pdf in (m, sample)
  //////////////////////////////////////////////

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

  // Associate models with states
  simPdf.addPdf(*massModelW,  "Wide");
  simPdf.addPdf(*massModelNW, "NoWide");


  
  /////////////////////
  // Get initial values
  /////////////////////
  RooArgSet *obs           = new RooArgSet(*m);
  RooArgSet *fitParameters = simPdf.getParameters(*obs);
  fitParameters->Print("v");
  cout << endl;

  
  
  ///////////////////////////
  // Perform simultaneous fit
  ///////////////////////////
  // Perform the fit
  RooAbsReal   *nll;
  RooMinuit    *minuit;
  RooFitResult *fitResult;
  if (!checkGeneration)
    {
      nll    = simPdf.createNLL(combData, Range("FitRange"), SplitRange(kTRUE), Verbose(kTRUE), NumCPU(cpus), Extended(kTRUE));
      minuit = new RooMinuit(*nll);
      
      minuit->setVerbose(kTRUE);
      minuit->setPrintLevel(3);

      minuit->migrad();
      minuit->hesse();
      fitResult = minuit->save();
      
      cout << endl
	   << "==================================="   << endl
	   << "=" << endl
	   << "= Status  = " << fitResult->status()   << endl
	   << "= Edm     = " << fitResult->edm()      << endl
	   << "= CovQual = " << fitResult->covQual()  << endl
	   << "=" << endl
	   << "==================================="   << endl
	   << endl;
      fitParameters->Print("v");
    }


  
  ///////////////
  // Plot results
  ///////////////

  // Wide
  RooPlot* frameW = m->frame(Title("Wide sample"), Range("FitRange_Wide"));
  binnedW->plotOn(frameW);
  // combData.plotOn(frameW,  Cut("sample==sample::Wide"));
  if (!checkGeneration)
    {
      simPdf.plotOn(  frameW,  Slice(sample, "Wide"),   ProjWData(sample, combData));
      simPdf.plotOn(  frameW,  Slice(sample, "Wide"),   Components("pdfComb"), ProjWData(sample, combData), LineStyle(kDashed), LineColor(kRed));
    }
  
  // NoWide
  RooPlot* frameNW = m->frame(Title("NoWide sample"), Range("FitRange_Wide"));
  binnedNW->plotOn(frameNW);
  // combData.plotOn(frameNW, Cut("sample==sample::NoWide"));
  if (!checkGeneration)
    {
      simPdf.plotOn(  frameNW, Slice(sample, "NoWide"), ProjWData(sample, combData));
      simPdf.plotOn(  frameNW, Slice(sample, "NoWide"), Components("pdfComb"), ProjWData(sample, combData), LineStyle(kDashed), LineColor(kRed));
    }
  
  // Draw and save canvases
  TCanvas* cW = new TCanvas("CanvasWide", "Wide", 0, 0, 900, 900);
  cW->cd();
  gPad->SetLeftMargin(0.15);
  frameW->GetYaxis()->SetTitleOffset(1.4);
  frameW->Draw();
  cW->Update();
  // cW->SaveAs(Form("OutputFiles/Bin%04d/%d.%s.%s.Bin%04d.Wide.FitYieldsSimultaneously.Tail%d.png",      bin, year, magnet.Data(), charge.Data(), bin, tailType));
  cW->SetLogy();
  cW->Update();
  // cW->SaveAs(Form("OutputFiles/Bin%04d/%d.%s.%s.Bin%04d.Wide.FitYieldsSimultaneously.Tail%d.Logy.png", bin, year, magnet.Data(), charge.Data(), bin, tailType));

  TCanvas* cNW = new TCanvas("CanvasNoWide", "NoWide", 1000, 0, 900, 900);
  cNW->cd();
  gPad->SetLeftMargin(0.15);
  frameNW->GetYaxis()->SetTitleOffset(1.4);
  frameNW->Draw();
  cNW->SetLogy();
  cNW->Update();

  if (!checkGeneration)
    {
      printf("\nGenerated vs fit:\n");
      printf("   signalW  = %8d ==> %10.2f +- %10.2f\n", genSigW,   signalW ->getVal(), signalW ->getError());
      printf("   combW    = %8d ==> %10.2f +- %10.2f\n", genCombW,  combW   ->getVal(), combW   ->getError());
      printf("   signalNW = %8d ==> %10.2f +- %10.2f\n", genSigNW,  signalNW->getVal(), signalNW->getError());
      printf("   combNW   = %8d ==> %10.2f +- %10.2f\n", genCombNW, combNW  ->getVal(), combNW  ->getError());
      
      cout << endl;
    }
  
  return;
}