using namespace RooFit;

void Minimal_WE()
{
  //Import the selected data
  TString filename("root://eoslhcb.cern.ch//eos/lhcb/user/c/chernov/public/KstGamma_MC12_selected.root");
  TFile* input_file = new TFile(filename,"READ");

  RooRealVar DMass("DMass","DMass",1680.,2080.);
  RooRealVar DeltaM("#Delta M", "#Delta M",139.57,152.); 
  RooRealVar Hel("cos(#theta)","cos(#theta)",-1.,1.);

  RooArgSet vars = RooArgSet(DMass,DeltaM,Hel);
  RooDataSet* _arr = (RooDataSet*)input_file->Get("_arr");

  //Create a model
  ////Delta M
  //RooRealVar lambda("#lambda","scale factor",1.0,0.8,4.0);
  RooRealVar lp0("lp0","constant term",266.3);
  RooRealVar lp1("lp1","linear term",-142.94);
  RooRealVar lp2("lp2","quadratic term",-75.41,-90,-60); lp2.setConstant(kTRUE);
  RooRealVar lp3("lp3","cubic term",40.58,33.3,47.5); lp3.setConstant(kTRUE);
  RooFormulaVar lambda("#lambda","@1+@2*(@0/1000)+@3*((@0/1000)**2)+@4*((@0/1000)**3)",RooArgList(DMass,lp0,lp1,lp2,lp3) );

  RooRealVar gmean_core_dm("#mu_{core}[#Delta M]","mean of the core part",145.436,143.0,147.); //gmean_core_dm.setConstant(kTRUE);
  RooRealVar gsigma1_dm_base("#sigma_{nrw;BASE}","the narrow sigma",.1,.8); gsigma1_dm_base.setVal(0.48); //gsigma1_dm_base.setConstant(kTRUE);
  RooFormulaVar gsigma1_dm("#sigma_{nrw}","@0*@1",RooArgList(gsigma1_dm_base,lambda) );
  RooGaussian g1_dm("G_{nrw}","Narrow Gaussian in DM",DeltaM,gmean_core_dm,gsigma1_dm);

  RooRealVar rho_s2_s1("#rho(#sigma_{2}/#sigma_{1})","ration between the broad and the narrow widths",1.83,1.4,2.2); //rho_s2_s1.setConstant(kTRUE);
  RooRealVar gsigma2_dm_base("#sigma_{brd;BASE}","the broad sigma",.7,1.4); gsigma2_dm_base.setVal(0.875); //gsigma2_dm_base.setConstant(kTRUE);
  RooFormulaVar gsigma2_dm("#sigma_{brd}","@0*@1*@2",RooArgList(gsigma1_dm_base,rho_s2_s1,lambda) );
  RooGaussian g2_dm("G_{brd}","Broad Gaussian in DM",DeltaM,gmean_core_dm,gsigma2_dm);
  
  RooRealVar gmean_tail_dm("#mu_{tail}[#Delta M]","mean of the third Gaussian",146.,144.0,148.0); gmean_tail_dm.setConstant(kTRUE);
  RooRealVar gsigma3_dm("#sigma_{tail}[#Delta M]","width of the tail",1.92,1.0,5.0); gsigma3_dm.setVal(1.92); gsigma3_dm.setConstant(kTRUE);

  RooGaussian g3_dm("G_{tail}","Tails Gaussian",DeltaM,gmean_tail_dm,gsigma3_dm);

  //recursive fractions
  RooRealVar f_nrw("f_{nrw}","fraction of a narrow Gaussian in a core double",0.61,0.015,0.985); //f_nrw.setConstant(kTRUE);
  RooAddPdf Core_Gaussian = RooAddPdf("Core Gaussian","Core PDF - sum of Gaussians",RooArgList(g1_dm,g2_dm),RooArgList(f_nrw) );
  RooRealVar f_core("f_{core}","What percentage of PDF is the core part",0.97,0.75,0.999); //f_core.setConstant(kTRUE);
  RooAddPdf DM_total("DM_total","Total PDF description of #Delta M peak",RooArgList(Core_Gaussian, g3_dm),RooArgList(f_core) );

  //M(D0)
  RooRealVar gmean_core_m("#mu_{core}[M]}","mean of the core Gaussian",1857.4,1820,1880); //gmean_core_m.setConstant(kTRUE);
  RooRealVar gsigma1_m("#sigma_{nrw}[M]","narrowest part of the mass model",26.4,15,30); //gsigma1_m.setConstant(kTRUE); 
  RooGaussian g1_m("G_{1}","Narrow Gaussian in M",DMass,gmean_core_m,gsigma1_m); 
  RooRealVar gmean_offset("#delta #mu_{brd}","offset of the broad Gaussian wrt narrow",32,-40,40); //gmean_offset.setConstant(kTRUE); 
  RooRealVar gsigma2_m("#sigma_{brd}[M]","broad part of the signal model",33.6,30,60.); //gsigma2_m.setConstant(kTRUE);
  RooFormulaVar gmean2_core_m("#mu_{2}[M]","@0+@1",RooArgList(gmean_core_m,gmean_offset));
  RooGaussian g2_m("G_{2}","Broad Gaussian in M",DMass,gmean2_core_m,gsigma2_m);

  //tail
  RooRealVar gmean_tail_m("#mu_{tail}[M]","",1884,1870,1920); gmean_tail_m.setConstant(kTRUE); 
  RooRealVar gsigma3_m("#sigma_{tail}[M]","",102,60,120); gsigma3_m.setConstant(kTRUE); 
  RooGaussian g3_m("G_{3}","",DMass,gmean_tail_m,gsigma3_m);
  
  RooRealVar f_core_m("f_{core}^{M}","",0.945,0.9,0.999); //f_core_m.setConstant(kTRUE);
  RooRealVar f_gm("f_{1}^{M}","fraction of the narrow Gaussian",0.82,0.5,0.95); //f_gm.setConstant(kTRUE);
  
  RooAddPdf M_DG("M_DG","",g1_m,g2_m,f_gm);
  RooAddPdf M_total("M_total","",M_DG,g3_m,f_core_m);

  //helicity
  //Sig HEL fixed from signal MC fit result.
  RooRealVar a1_hel("a1_hel", "Linear part of signal", 0.152, -5., 5.); a1_hel.setConstant(kTRUE);
  RooRealVar a2_hel("a2_hel", "Quadratic part of signal acceptance", -0.313, -5., 5.); a2_hel.setConstant(kTRUE);
  RooGenericPdf Hel_sig_corrected("Hel_signal", "(1. - (@0*@0))  * (1 + @1 * @0 + @2 * @0 * @0)", RooArgList(Hel,  a1_hel, a2_hel));
  RooGenericPdf Hel_sig_theory("Hel_signal","(1- (@0**2) )",RooArgList(Hel));

  //non-commutativity between addition and multiplication?

  //case 1: (this seems to be working)
  /*
  RooProdPdf core_1("core_1","",RooArgList(M_total,Hel_sig_corrected),Conditional(g1_dm, DeltaM ) );
  RooProdPdf core_2("core_2","",RooArgList(M_total,Hel_sig_corrected),Conditional(g2_dm, DeltaM ) );
  RooProdPdf core_3("core_3","",RooArgList(M_total,Hel_sig_corrected),Conditional(g3_dm, DeltaM ) );

  RooAddPdf correlated_3d("correlated_3d","",core_1,core_2,f_nrw);
  RooAddPdf total("total","",correlated_3d,core_3,f_core);
  */

  //case 2:
  RooProdPdf total("total","overall 3D PDF",RooArgList(M_total,Hel_sig_corrected),Conditional(DM_total, RooArgSet(DeltaM) ) );

  total.fitTo(*_arr);

  //plotting
  TCanvas plots;
  const Int_t w = 1600; const Int_t h = 800;
  plots.SetCanvasSize(w,h);
  plots.Divide(3,1);

  plots.cd();
  TVirtualPad* cpad;
  cpad = plots.GetPad(1);
  cpad->Divide(1,2);
  cpad->GetPad(1)->SetPad(0.05,0.225,0.95,0.975);                                                                                                                                                        
  cpad->GetPad(2)->SetPad(0.05,0.025,0.95,0.225);                                                                                                                                                        
  cpad->cd(1); 
  RooPlot* france = DMass.frame(1680,2080,100); france->SetTitle("M(D^{0})");
  _arr->plotOn(france); total.plotOn(france, LineColor(44), LineStyle(kDashed), Components("G_{tail}")); total.plotOn(france, LineColor(14), LineStyle(kDashed), Components("G_{2}") ); total.plotOn(france, LineColor(kBlack) ); total.paramOn(france, Layout(0.6,0.9,0.9), Format("NE",AutoPrecision(2) ), Parameters(RooArgSet(gmean_core_m, gsigma1_m,gmean_offset,gsigma2_m,gmean_tail_m,gsigma3_m,f_core_m,f_gm) ) );
  RooHist* m_pulls = france->pullHist();
  RooPlot* france_pulls = DMass.frame(1680,2080,100); france_pulls->addPlotable(m_pulls, "P");
  france->Draw();
  cpad->cd(2);
  france_pulls->SetTitle(""); france_pulls->Draw("");

  cpad = plots.GetPad(2);
  cpad->Divide(1,2);
  cpad->GetPad(1)->SetPad(0.05,0.225,0.95,0.975);
  cpad->GetPad(2)->SetPad(0.05,0.025,0.95,0.225);
  cpad->cd(1);
  RooPlot* austria = Hel.frame(-1.0,1.0,100); austria->SetTitle("cos(#theta)");
  _arr->plotOn(austria); total.plotOn(austria, LineColor(kBlack) ); total.paramOn(austria, Layout(0.6,0.9,0.9), Format("NE", AutoPrecision(2) ), Parameters(RooArgSet(lp3,lp2) ) );
  RooHist* hel_pulls = austria->pullHist();
  RooPlot* austria_pulls = Hel.frame(-1,1,100); austria_pulls->addPlotable(hel_pulls,"P");
  austria->Draw();
  cpad->cd(2);
  austria_pulls->SetTitle(""); austria_pulls->Draw("");

  cpad = plots.GetPad(3);
  cpad->Divide(1,2);
  cpad->GetPad(1)->SetPad(0.05,0.225,0.95,0.975);
  cpad->GetPad(2)->SetPad(0.05,0.025,0.95,0.225);
  cpad->cd(1);
  RooPlot* serbia = DeltaM.frame(139.57,152.,50); serbia->SetTitle("#Delta M");
  _arr->plotOn(serbia); total.plotOn(serbia, LineColor(kBlack) );   RooHist* dm_pulls = serbia->pullHist(); total.paramOn(serbia, Layout(0.6,0.9,0.9), Format("NE",AutoPrecision(2) ), Parameters(RooArgSet(gmean_core_dm,gmean_tail_dm,gsigma1_dm_base,rho_s2_s1) ) );
  RooPlot* serbia_pulls = DeltaM.frame(139.57,152.,50); serbia_pulls->addPlotable(dm_pulls,"P");
  serbia->Draw();
  cpad->cd(2);
  serbia_pulls->SetTitle(""); serbia_pulls->Draw("");

  plots.SaveAs("test.pdf");

}