// ROOT
#include "TCanvas.h"

// RooFit - Basics
#include "RooArgList.h"
#include "RooArgSet.h"
#include "RooCategory.h"
#include "RooDataSet.h"
#include "RooFormulaVar.h"
#include "RooPlot.h"
#include "RooRealVar.h"

// RooFit - Pdf
#include "RooLognormal.h"

// RooFit - Pdf AbsAnaConv
#include "RooBDecay.h"

// RooFit - Resolution
#include "RooGaussModel.h"

using namespace RooFit;

void CondTest() {
  bool use_per_event_resol = true;
  int  num_cpu = 32;

  // observables
  RooRealVar   t("t","t",0.,13.,"ps");
  RooRealVar   terr("terr","#sigma_{t}",1e-5,0.2,"ps");
  RooCategory  tag("tag","d");
  tag.defineType("Pos", +1);
  tag.defineType("Neg", -1);



  // per-event error pdf
  RooRealVar    parTerrMedian("parTerrMedian","m_{#sigma_{t}}",0.03);
  RooRealVar    parTerrSigma("parTerrSigma","#sigma_{#sigma_{t}}",0.6);
  RooLognormal  pdfTerr("pdfTerr","P(#sigma_{t})",terr,parTerrMedian,parTerrSigma);

  // resolution
  RooResolutionModel* resTime = NULL;

  RooRealVar parTimeResBias("parTimeResBias","b_{t}",0.,"ps");
  RooRealVar parTimeResScale("parTimeResScale","S",1.);

  if (!use_per_event_resol) {
    resTime = new RooGaussModel("resTime","R(t-t')",t,parTimeResBias,parTerrMedian);
  } else {
    resTime = new RooGaussModel("resTime","R(t-t')",t,parTimeResBias,terr,parTimeResScale);
  }

  // t pdf
  RooRealVar Tau("Tau","#tau",1.5,"ps"); 
  RooRealVar Dm("Dm","#Delta m",1.,"ps^{-1}");
  RooRealVar DG("DG","#Delta#Gamma",0.);

  RooRealVar    Chf("Chf","Ch_{f}",1.0);
  RooRealVar    Shf("Shf","Sh_{f}",0.0);
  RooRealVar    Cf("Cf","C_{f}",0.0);
  RooFormulaVar Sf("Sf","-1.*@0",RooArgList(tag));

  RooBDecay pdfTime("pdfTime","pdfTime",t,Tau,DG,Chf,Shf,Cf,Sf,Dm,*resTime,RooBDecay::SingleSided);

  // full pdf
  RooAbsPdf* pdfFull = NULL;

  if (!use_per_event_resol) {
    pdfFull = new RooProdPdf("pdfFull","P(t,d,#sigma_{t})",RooArgSet(pdfTerr,pdfTime));
  } 
  else {
    pdfFull = new RooProdPdf("pdfFull","P(t,d,#sigma_{t})",RooArgSet(pdfTerr),Conditional(RooArgSet(pdfTime),RooArgSet(t,tag)));
  }

  // generate dataset
  RooDataSet* data = pdfFull->generate(RooArgSet(t,tag,terr),100000);


  // plot
  TCanvas c("c","c",800,600);

  RooPlot* plot = NULL;

  plot = t.frame();
  data->plotOn(plot);
  pdfFull->plotOn(plot);
  plot->Draw();
  c.SaveAs("Time.png");
  delete plot;

  plot = terr.frame();
  data->plotOn(plot);
  pdfFull->plotOn(plot);
  plot->Draw();
  c.SaveAs("TimeErr.png");
  delete plot;

  plot = t.frame();
  data->plotOn(plot,Asymmetry(tag));
  pdfFull->plotOn(plot,Asymmetry(tag));
  plot->Draw();
  c.SaveAs("Asymmetry.png");
  delete plot;

  plot = t.frame();
  data->plotOn(plot,Asymmetry(tag));
  pdfFull->plotOn(plot,Asymmetry(tag),ProjWData(RooArgSet(terr),*data,kTRUE),NumCPU(num_cpu));
  plot->Draw();
  c.SaveAs("AsymmetryProjWData.png");
  delete plot;

  plot = t.frame();
  data->plotOn(plot,Cut("tag==+1"));
  pdfFull->plotOn(plot,NumCPU(num_cpu),Slice(tag,"Pos"));
  plot->Draw();
  c.SaveAs("TimeTagPos.png");
  delete plot;

  plot = t.frame();
  data->plotOn(plot,Cut("tag==-1"));
  pdfFull->plotOn(plot,NumCPU(num_cpu),Slice(tag,"Neg"));
  plot->Draw();
  c.SaveAs("TimeTagNeg.png");
  delete plot;


  delete data;
  delete pdfFull;
  delete resTime;
}