// SS 02/14/12

// Construct a conditional PDF ("smear") with an energy-dependant
// energy resolution by taking the product of a physics "model"
// and a gaussian with varying sigma.
// The sigma depends on a dummy energy, "EP".
// Construct a data set with observable energy, "E", and dummy
// energy "EP", and use the data set to project out EP to
// do a fit to E. 
// The resolution model will be G(E, mu=EP, sigma=sigma(EP))

// SS 02/17/12
// Update the RooProdPdf to be a conditional product, as per
// WV's suggestion: http://root.cern.ch/phpBB3/viewtopic.php?f=15&t=14160
//  2.1: try treating the new pdf as non-conditional 

// APPEARS TO FIT PROPERLY THE SIGNAL FRACTIONS... BEST SCRIPT TO DATE

{
  using namespace RooFit;
  gSystem->Load("libRooFit");
  
  //// Create PDFs

  // Create Observables
  RooRealVar E("E","E", 0., 4000.); // physical energy variable
  RooRealVar EP("EP","EP", 0., 4000.); // dummy energy variable
  E.setBins(1000.);
  EP.setBins(1000.);
  
  // G ( E | EP )
  // Define an energy-dependent resolution function
  RooRealVar a0("a0","a0", 24.); // define parameters for sigma
  RooRealVar a1("a1","a1",0.3); // " "        " " 
  RooRealVar a2("a2","a2",0.04, 0.0, 1.0); // " "        " "
  RooFormulaVar gsigma("gsigma","@0 + @1*sqrt(@3) + @2*@3", RooArgList(a0, a1, a2, EP));
  
  // Create Gaussian resolution model g(E, EP, sigma(EP))
  RooGaussian gm("gm","gauss model with varying sigma", E, EP, gsigma);

  // Plot sigma(EP)
  a2.setVal(0.08); // vary one of the parameters to make the fit interesting
  RooPlot * test = EP.frame(Title("Resolution Function"));
  test->SetXTitle("Dummy Energy");
  test->SetYTitle("Resolution");
  gsigma.plotOn(test, LineColor(kRed));//plot in red before the fit (a2_init=0.08, large)
  
  // F ( EP )
  // "Model" PDF: corresponds to our physics + background 
  // Just 3 sharp gaussian
  RooGaussian mgaus1("mgaus1", "mgaus1", EP, RooConst(600.), RooConst(3.0));
  RooGaussian mgaus2("mgaus2", "mgaus2", EP, RooConst(1600.), RooConst(3.0));
  RooGaussian mgaus3("mgaus3", "mgaus3", EP, RooConst(2600.), RooConst(3.0));
  RooRealVar fsig1("fsig1", "fsig1", 0.13, 0., 1.);
  RooRealVar fsig2("fsig2", "fsig2", 0.33, 0., 1.);
  RooAddPdf model("model","model",RooArgList(mgaus1, mgaus2, mgaus3), RooArgList(fsig1, fsig2));

  // plot the model
  RooPlot * modelplot = EP.frame(Title("Physics Model PDF (Unsmeared)"));
  model.plotOn(modelplot);

  // F( EP ) x G ( E | EP )
  //  RooProdPdf smear("smear","smear", RooArgSet(gm, model) );
  // Make a conditional product
  // F(X|Y)*G(Y) = RooProdPdf prod("prod","prod",G,Conditional(F,x));
  RooProdPdf smear("smear","smear",model,Conditional(gm, E));

  // Generate a fake "measured" distribution, including both E and EP (same for every point) 
  // by hand following example in the root tutorial /roofit/rf303_conditional.C 
  // Have two columns, one with E and one with EP.
  
  RooArgSet coord(E,EP) ;
  RooDataSet* fakeDataEEP = new RooDataSet("fakeDataEEP","fakeDataEEP",RooArgSet(E,EP)) ;
  // Let EP be distributed in 3 "delta functions" at given energies
  // Then let E be that energy smeared according to a gaussian
  for (int i=0 ; i<10000 ; i++) {
    EP = 600.;
    Double_t tmpE = gRandom->Gaus(600., 55.); // 24 + 0.3*sqrt(600) + 600*0.04 = 55.3
    if (fabs(tmpE)< 3000. && tmpE > 0.0) { 
      E = tmpE ;  
      fakeDataEEP->add(coord) ;
    }
    if (i < 5000){
      EP = 1600.;
      tmpE = gRandom->Gaus(1600., 100.); // 24 + 0.3*sqrt(1600) + 1600*0.04 = 100.
      if (fabs(tmpE)< 3000. && tmpE > 0.0) { 
	E = tmpE ;  
	fakeDataEEP->add(coord) ;
      }
    }
    EP = 2600.;
    tmpE = gRandom->Gaus(2600., 143.); // 24 + 0.3*sqrt(2600) + 2600*0.04 = 143.
    if (fabs(tmpE)< 3000. && tmpE > 0.0) { 
      E = tmpE ;  
      fakeDataEEP->add(coord) ;
    }
  }
  
  // Make a subset of the fake data only containing EP -- still following rf303
  RooDataSet* fakeDataEP = (RooDataSet*) fakeDataEEP->reduce(EP) ;
  
  // Plot data, and plot the "smear" PDF before doing the fit
  RooPlot * dataplot = E.frame(Title("Fake Data Set and Fit to Smeared PDF (Projected)"));
  fakeDataEEP->plotOn(dataplot);
  smear.plotOn(dataplot, ProjWData(*fakeDataEP, kTRUE), LineColor(kRed));

  // Do a fit of the smeared PDF to full data set, specifying EP as conditional
  // Need E and EP info to do a fit, so use full dataset
  //  smear.fitTo(*fakeDataEEP, ConditionalObservables(RooArgList(EP)), NumCPU(2));
  smear.fitTo(*fakeDataEEP, NumCPU(2));

  // Now plot the "smear" PDF after the fit.
  smear.plotOn(dataplot, ProjWData(*fakeDataEP, kTRUE));
  smear.paramOn(dataplot, Layout(0.65));
  gsigma.plotOn(test);//plot after the fit in blue, same as model
  
  // Make two-dimensional plot of conditional p.d.f in (E, EP)
  TH1* hh_smear = smear.createHistogram("hh_smear",E,Binning(400),YVar(EP,Binning(1000))) ; 
  hh_smear->SetTitle("Two-Dimensional Smeared PDF"):

  TH2F* hh_data = fakeDataEEP.createHistogram(E, EP, 400, 1000, "","hh_data");

  TCanvas * c = new TCanvas("resSmear", "resSmear", 1200, 800);
  c->Divide(2,2);
  c->cd(1); test->Draw();
  c->cd(2); modelplot->Draw();
  c->cd(3); dataplot->Draw();
  c->cd(4);
  gStyle->SetPalette(1);
  hh_smear->Draw("colz");

  smear.Print("V");
  return; 
}