#include "RooRealVar.h"
#include "RooBinning.h"
#include "RooDataSet.h"
#include "RooWorkspace.h"
#include "RooPlot.h"
#include "RooBreitWigner.h"
#include "RooFFTConvPdf.h"
#include "RooFitResult.h"

#include "myRooPdfs.h"

#include "TFile.h"
#include "TF1.h"
#include "TTree.h"
#include "TChain.h"
#include "TMath.h"
#include "TROOT.h"
#include "TCanvas.h"
#include "TAxis.h"

#include <fstream>

using namespace RooFit;

// ============================================
// to be modified:
//Int_t MINmass= 1250;
//Int_t MAXmass= 1700;
Int_t MINmass= 300;
Int_t MAXmass= 6000;
// ============================================

RooArgSet* defineVariables() {

  RooRealVar* mgg        = new RooRealVar("mgg",        "M(gg)",       MINmass, MAXmass, "GeV");
  RooRealVar* eventClass = new RooRealVar("eventClass", "eventClass",    -10,      10,   "");
  RooRealVar* weight     = new RooRealVar("weight",     "weightings",      0,     1000,  "");

  RooArgSet* ntplVars = new RooArgSet(*mgg, *eventClass, *weight);
  return ntplVars;
}

// loading signal data and making roodatasets   
void AddSigData(RooWorkspace* w) {

  // Variables       
  RooArgSet* ntplVars = defineVariables();
  
  // Input files          
  TChain* sigTree = new TChain();
  sigTree->Add("/afs/cern.ch/user/c/crovelli/public/testFile1500.root/DiPhotonTree");
  
  TString mainCut = TString::Format("mgg>=300 && mgg<=6000 && eventClass==0");
  RooDataSet sigWeighted("sigWeighted","dataset",sigTree,*ntplVars,mainCut,"weight");
  w->import(sigWeighted,Rename("SigWeight_cat0"));
}
    
void MakeIT(RooWorkspace* w){
  
  TString wsDir = "workspaces/";
  RooWorkspace *wAll = new RooWorkspace("w_all","w_all");  
  
  RooRealVar *mgg = new RooRealVar("mgg", "mgg", 1250, 1700);

  // dataset
  wAll->import(*w->data("SigWeight_cat0"));   
  
  // convolution : BW 
  RooRealVar BWmean("BWmean_cat0","",1500);
  RooRealVar BWwidth("BWwidth_cat0","",21);
  RooBreitWigner SigModelBW("SigModelBW_cat0","SigModelBW_cat0", *mgg, BWmean, BWwidth);

  // convolution : dCB 
  RooRealVar CBsigma("CBsigma_cat0","",9.20534);
  RooRealVar CBalpha1("CBalpha1_cat0","",0.858349);
  RooRealVar CBalpha2("CBalpha2_cat0","",1.70584);
  RooRealVar CBmean("CBmean_cat0","",-4.24842);
  RooRealVar CBn1("CBn1_cat0","",3.79562);
  RooRealVar CBn2("CBn2_cat0","",3.56032);
  RooDoubleCB ResponseDoubleCB("ResponseDoubleCB_cat0","ResponseDoubleCB_cat0", *mgg, CBmean, CBsigma, CBalpha1, CBn1, CBalpha2, CBn2);
    
  // convolution
  mgg->setBins(5000, "cache");  
  mgg->setBinning(RooBinning(5000,1250,1700));
  
  RooFFTConvPdf* ConvolutedRes;
  ConvolutedRes = new RooFFTConvPdf("mggSig_cat0","mggSig_cat0", *mgg,SigModelBW, ResponseDoubleCB);
  
  // dummy fit to fix the binning
  RooDataSet *sigToFit = (RooDataSet*) w->data("SigWeight_cat0");
  RooFitResult* fitresults_CB = (RooFitResult* ) ConvolutedRes->fitTo(*sigToFit, SumW2Error(kFALSE), Range(1250,1700), RooFit::Save(kTRUE));
  wAll->import(*ConvolutedRes);

  // control plot
  RooPlot* controlPlot  = w->var("mgg")->frame(Range(1250, 1700),Title(""),Bins(60));
  controlPlot->SetTitle("");
  sigToFit->plotOn(controlPlot);
  float max = controlPlot->GetMaximum();
  SigModelBW.plotOn(controlPlot, LineColor(kGreen), LineStyle(kDashed));
  ConvolutedRes->plotOn(controlPlot, LineColor(kBlue));  
  TCanvas* c1 = new TCanvas("c1","PhotonsMass",0,0,800,800);
  c1->cd(1);
  controlPlot->Draw();  
  controlPlot->GetYaxis()->SetRangeUser(0.01, max*1.2);
  c1->SaveAs("closureTest.png");

  // saving
  TFile fileWs("dummy.root","RECREATE");
  fileWs.cd();
  wAll->writeToFile("dummy.root");
  
  return;
}

void runAllParametric(){
  
  RooWorkspace *w = new RooWorkspace("w","w");  
  AddSigData(w);   
  MakeIT(w);
  
  return;
}