#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooStats/SPlot.h"
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooChebychev.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooAddition.h"
#include "RooProduct.h"
#include "TCanvas.h"
#include "RooAbsPdf.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooWorkspace.h"
#include "RooConstVar.h"

// use this order for safety on library loading
using namespace RooFit ;
using namespace RooStats ;


// see below for implementation
void AddModel(RooWorkspace*);
void AddData(RooWorkspace*);
void DoSPlot(RooWorkspace*);
void MakePlots(RooWorkspace*);

void plot()
{

  // Create a new workspace to manage the project.
  RooWorkspace* wspace = new RooWorkspace("myWS");

  // add the signal and background models to the workspace.
  // Inside this function you will find a discription our model.
  AddModel(wspace);

  // make our canvas
  TCanvas* cdata = new TCanvas("sPlot","sPlot demo", 400, 600);
  cdata->Divide(2,2);

  // add some toy data to the workspace
  AddData(wspace);

  // inspect the workspace if you wish
  //  wspace->Print();

  // do sPlot.  
  //This wil make a new dataset with sWeights added for every event.
  DoSPlot(wspace,cdata);

  // Make some plots showing the discriminating variable and 
  // the control variable after unfolding.
  MakePlots(wspace,cdata);

  // cleanup
  delete wspace;
  
}

 
//____________________________________
void AddModel(RooWorkspace* ws)
{
    RooRealVar Mass("Mass","invariant mass (GeV/c2)",1,5);
      
    //Breit-Wigner for the phi
    
    RooRealVar mean_phi("mean_phi","mean_phi",1.0,0.9,1.1);
    RooRealVar g("sigma_phi","sigma_phi",1,0,2);
    RooBreitWigner *phi_BW = new RooBreitWigner("phi_BW","Breit Wigner PDF",Mass,mean_phi,g);
    
    //Crystal ball for the J/psi
    
    RooRealVar mean_jpsi("mean_jpsi","mean_jpsi",3,2.9,3.2);
    RooRealVar sigma_jpsi("sigma_jpsi","sigma_jpsi",0.07,0,0.1);
    RooRealVar alpha("alpha","alpha",1);
    RooRealVar n("n","n",5);
    RooCBShape *jpsi = new RooCBShape("jpsi","crystal ball PDF",Mass,mean_jpsi,sigma_jpsi,alpha,n);
    
    //Crystal ball for the psi
    
    RooRealVar mean_psi("mean_psi","mean_psi",3.6,3.4,3.8);
    RooRealVar sigma_psi("sigma_psi","sigma_psi",0.07,0,0.1);
    RooRealVar alpha_psi("alpha","alpha",1.5);
    RooRealVar n_psi("n","n",1);
    RooCBShape *psi = new RooCBShape("psi","crystal ball PDF",Mass,mean_psi,sigma_psi,alpha_psi,n_psi);
    
    //Crystal ball for the Upsilon
    
    RooRealVar mean_up("mean_up","mean_up",9.4,9.3,9.5);
    RooRealVar sigma_up("sigma_up","sigma_up",0.07,0,0.1);
    RooRealVar alpha_up("alpha_up","alpha_up",1,0,5);
    RooRealVar n_up("n_up","n_up",1,0,10);
    RooCBShape *upsilon = new RooCBShape("upsilon","crystal ball PDF",Mass,mean_up,sigma_up,alpha_up,n_up);
    
    //Exponential background
    RooRealVar a("a","a",-0.85,-0.9,0);
    RooExponential bkg("exp","exp",Mass,a);
    
    RooRealVar fsig("fsig","signal",200.0,0.,1e7.);
    RooRealVar fsig2("fsig2","signal2",200.0,0.,1e7);
    RooRealVar fsig3("fsig3","signal3",0.1,0.,1.);
    RooRealVar fsig4("fsig4","signal4",0.1,0.,1.);
    RooRealVar fsig5("fsig5","signal5",0.1,0.,1.);
    
    RooAddPdf model("model","model",RooArgList(*jpsi,bkg),RooArgList(fsig,fsig2));

    ws->import(model);
}

//____________________________________
void AddData(RooWorkspace* ws){
  // Add a toy dataset

  // how many events do we want?
  Int_t nEvents = 30000;

  // get what we need out of the workspace to make toy data
  RooAbsPdf* model = ws->pdf("model");
  RooRealVar* invMass = ws->var("Mass");
 
  // make the toy data
  std::cout << "make data set and import to workspace" << std::endl;
  RooDataSet* data = model->generate(RooArgSet(*invMass),nEvents);
  
  // import data into workspace
  ws->import(*data, Rename("data"));
    
    /*RooRealVar *Mass = ws->var("Mass");
    RooRealVar trackpt("trackpt","trackpt",-1e200,1e200,"GeV");
    RooRealVar tracketa("tracketa","tracketa",-1e200,1e200,"");
    RooRealVar trackphi("trackphi","trackphi",-1e200,1e200,"");
    RooRealVar cl_E("cl_E","cl_E",-1e200,1e200,"GeV");
    RooRealVar Eratio("Eratio","E_{ratio}",-1e200,1e200,"");
    RooRealVar ET("ET","ET",-1e200,1e200,"");
    RooRealVar cl_eta("cl_eta","cl_eta",-1e200,1e200,"");
    RooRealVar d0Signfcns("sgnfcns_trackd0pv","sgnfcns_trackd0pv" ,-1e200,1e200,"");
    RooRealVar depth("depth","depth",-1e200,1e200,"");
    RooRealVar f1("f1","f1",-1e200,1e200,"");
    RooRealVar f1core("f1core","f1core" ,-1e200,1e200,"");
    RooRealVar f3("f3","f3" ,-1e200,1e200,"");
    RooRealVar f3core("f3core","f3core" ,-1e200,1e200,"");
    RooRealVar deltaEta1("deltaEta1","deltaEta1" ,-1e200,1e200,"");
    RooRealVar deltaPhiRescaled("deltaPhiRescaled","deltaPhiRescaled" ,-1e200,1e200,"");
    RooRealVar reta("reta","reta" ,-1e200,1e200,"");
    RooRealVar rphi("rphi","rphi" ,-1e200,1e200,"");
    RooRealVar weta2("weta2","weta2" ,-1e200,1e200,"");
    RooRealVar ws3("ws3","ws3" ,-1e200,1e200,"");
    RooRealVar wstot("wstot","wstot" ,-1e200,1e200,"");
    RooRealVar rHad("rHad","rHad" ,-1e200,1e200,"");
    RooRealVar rHad1("rHad1","rHad1" ,-1e200,1e200,"");
    RooRealVar trtHitRatio("trtHitRatio","trtHitRatio" ,-1e200,1e200,"");
    RooRealVar trtOutlierRatio("trtOutlierRatio","trtOutlierRatio" ,-1e200,1e200,"");
    RooRealVar nBLHits("nBLHits","nBLHits" ,-1e200,1e200,"");
    RooRealVar nSCTHits("nSCTHits","nSCTHits" ,-1e200,1e200,"");
    RooRealVar nPixelHits("nPixelHits","nPixelHits" ,-1e200,1e200,"");
    RooRealVar nSiDeadSensors("nSiDeadSensors","nSiDeadSensors" ,-1e200,1e200,"");

    ///Get data from the file
    TFile * dataFile= new TFile("output.root");
    TTree * tree    = (TTree*)dataFile->Get("ProbeElectrons");
    RooArgSet probeset;

    probeset.add(*Mass);
    probeset.add(trackpt);
    probeset.add(tracketa);
    probeset.add(trackphi);
    probeset.add(cl_E);
    probeset.add(Eratio);
    probeset.add(ET);
    probeset.add(cl_eta);
    probeset.add(d0Signfcns);
    probeset.add(depth);
    probeset.add(f1);
    probeset.add(f1core);
    probeset.add(f3);
    probeset.add(f3core);
    probeset.add(deltaEta1);
    probeset.add(deltaPhiRescaled);
    probeset.add(reta);
    probeset.add(rphi);
    probeset.add(weta2);
    probeset.add(ws3);
    probeset.add(wstot);
    probeset.add(rHad);
    probeset.add(rHad1);
    probeset.add(trtHitRatio);
    probeset.add(trtOutlierRatio);
    probeset.add(nBLHits);
    probeset.add(nSCTHits);
    probeset.add(nPixelHits);
    probeset.add(nSiDeadSensors);

    RooDataSet* data = new RooDataSet("data","data",tree,probeset);
    ws->import(*data,Rename("data"));*/
}

//____________________________________
void DoSPlot(RooWorkspace* ws,TCanvas* cdata){
  std::cout << "Calculate sWeights" << std::endl;

  // get what we need out of the workspace to do the fit
  RooAbsPdf* model = ws->pdf("model");
  RooRealVar* fsig = ws->var("fsig");
  RooRealVar* fsig2 = ws->var("fsig2");
  
  RooDataSet* data = (RooDataSet*) ws->data("data");
  RooRealVar* mass = ws->var("Mass");

  // fit the model to the data.
  model->fitTo(*data, Range(2.5,3.5) );
  
  cdata->cd(1);
  RooPlot* mframe = mass->frame();

  data->plotOn(mframe);
  model->plotOn(mframe);
  model->plotOn(mframe,Components("exp"),LineStyle(kDashed));
  mframe->SetTitle("Mass Fit before splot");

  mframe->Draw();


  // The sPlot technique requires that we fix the parameters
  // of the model that are not yields after doing the fit.
  RooRealVar* jpsi_mean = ws->var("mean_jpsi");
  RooRealVar* sigma_jpsi= ws->var("sigma_jpsi");
  RooRealVar* alpha     = ws->var("alpha");
  RooRealVar* n         = ws->var("n");

  RooRealVar * a        = ws->var("a");

  jpsi_mean->setConstant();
  sigma_jpsi->setConstant();
  alpha->setConstant();
  n->setConstant();
  a->setConstant();

  RooMsgService::instance().setSilentMode(true);


  // Now we use the SPlot class to add SWeights to our data set
  // based on our model and our yield variables
  RooStats::SPlot* sData = new RooStats::SPlot("sData","An SPlot",
		            *data, model, RooArgList(*fsig,*fsig2) );

  // Check that our weights have the desired properties

  std::cout << "Check SWeights:" << std::endl;


  // import this new dataset with sWeights
 std::cout << "import new dataset with sWeights" << std::endl;
 ws->import(*data, Rename("dataWithSWeights"));


}

void MakePlots(RooWorkspace* ws,TCanvas *cdata){

  // Here we make plots of the discriminating variable (invMass) after the fit
  // and of the control variable (isolation) after unfolding with sPlot.
  std::cout << "make plots" << std::endl;


  // get what we need out of the workspace
  RooAbsPdf* model      = ws->pdf("model");
  RooAbsPdf* jpsiModel  = ws->pdf("jpsi");
  RooAbsPdf* qcdModel   = ws->pdf("exp");

  RooRealVar* invMass = ws->var("Mass");

  // note, we get the dataset with sWeights
  RooDataSet* data = (RooDataSet*) ws->data("dataWithSWeights");

  // this shouldn't be necessary, need to fix something with workspace
  // do this to set parameters back to their fitted values.
  //model->fitTo(*data, Extended() );

  //plot invMass for data with full model and individual componenets overlayed
  //  TCanvas* cdata = new TCanvas();
  cdata->cd(2);
  RooPlot* frame = invMass->frame() ; 
  data->plotOn(frame ) ; 
  model->plotOn(frame) ;   
  model->plotOn(frame,Components(*jpsiModel),LineStyle(kDashed), LineColor(kRed)) ;   
  model->plotOn(frame,Components(*qcdModel),LineStyle(kDashed),LineColor(kGreen)) ;   
    
  frame->SetTitle("Fit model after splot");
  frame->Draw() ;
 

  // Now use the sWeights to show isolation distribution for Z and QCD.  
  // The SPlot class can make this easier, but here we demonstrait in more
  // detail how the sWeights are used.  The SPlot class should make this 
  // very easy and needs some more development.

  // Plot isolation for Z component.  
  // Do this by plotting all events weighted by the sWeight for the Z component.
  // The SPlot class adds a new variable that has the name of the corresponding
  // yield + "_sw".
  cdata->cd(3);

  // create weightfed data set 
  RooDataSet * dataw_jpsi = new RooDataSet(data->GetName(),data->GetTitle(),data,*data->get(),0,"fsig_sw") ;

  RooPlot* frame2 = invMass->frame() ; 
  dataw_jpsi->plotOn(frame2, DataError(RooAbsData::SumW2) ) ; 
    
  frame2->SetTitle("Jpsi from sWeight applied dataset");
  frame2->Draw() ;

  // Plot isolation for QCD component.  
  // Eg. plot all events weighted by the sWeight for the QCD component.
  // The SPlot class adds a new variable that has the name of the corresponding
  // yield + "_sw".
  cdata->cd(4);
  RooDataSet * dataw_qcd = new RooDataSet(data->GetName(),data->GetTitle(),data,*data->get(),0,"fsig2_sw") ;
  RooPlot* frame4 = invMass->frame() ; 
  dataw_qcd->plotOn(frame4,DataError(RooAbsData::SumW2) ) ; 
  frame4->SetTitle("QCD from sweight applied dataset");
    
  frame4->Draw() ;


  //  cdata->SaveAs("SPlot.gif");

}