using namespace RooFit; 
using namespace RooStats; 

void lambda_ubfit_final()
{
   
   //SetBelle2Style() ;
   
   double M,InvM, p, p_px, p_py, p_pz, Pi_px, Pi_py, Pi_pz, Pi_mcPDG, p_E, Pi_E, Pi_dr, Pi_dz, p_dz,p_nCDCHits, Pi_nCDCHits, Pi_nSVDHits, p_nSVDHits, p_cosTheta, Pi_cosTheta;
   double Pi_SVDdEdx, p_SVDdEdx, Pi_p, p_p;
   double chiProb, isSignal, cosalpha, distance, Pi_pSVD, p_pSVD ;
   
   TChain *chain = new TChain("tree");  
   chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/test_lambda/unbinned/ok/unbinned/lambda_draw_unb_all.root");
  
   chain->SetBranchAddress("M",&M);
   chain->SetBranchAddress("InvM",&InvM);            // add
   //chain->SetBranchAddress("p",& p);
   chain->SetBranchAddress("p_px",&p_px);
   chain->SetBranchAddress("p_py",&p_py);
   chain->SetBranchAddress("p_pz",&p_pz);
   chain->SetBranchAddress("Pi_px",&Pi_px);
   chain->SetBranchAddress("Pi_py",&Pi_py);
   chain->SetBranchAddress("Pi_pz",&Pi_pz);
   chain->SetBranchAddress("p_E",&p_E);
   chain->SetBranchAddress("Pi_E",&Pi_E);   
   chain->SetBranchAddress("distance",&distance);
   chain->SetBranchAddress("cosalpha",&cosalpha);
   chain->SetBranchAddress("p_cosTheta",&p_cosTheta);   
   chain->SetBranchAddress("Pi_pSVD",&Pi_pSVD);         // add
   chain->SetBranchAddress("p_pSVD",&p_pSVD);           // add
   chain->SetBranchAddress("Pi_SVDdEdx",&Pi_SVDdEdx);  // add
   chain->SetBranchAddress("p_SVDdEdx",&p_SVDdEdx);    // add
   chain->SetBranchAddress("Pi_p",&Pi_p);
   chain->SetBranchAddress("p_p",&p_p);   
   chain->SetBranchAddress("chiProb",&chiProb);
   chain->SetBranchAddress("isSignal",&isSignal);     // add
   chain->SetBranchAddress("p_nCDCHits",&p_nCDCHits);
   chain->SetBranchAddress("Pi_nCDCHits",&Pi_nCDCHits);  
   chain->SetBranchAddress("p_nSVDHits",&p_nSVDHits);   //add
   chain->SetBranchAddress("Pi_nSVDHits",&Pi_nSVDHits);  //add
   
   RooRealVar pP_over_pL("pP_over_pL","pP_over_pL",0.0,20.0);
   RooRealVar invm("invm","M_{p^{+}#pi^{-}}",1.1,1.13,"GeV/c^{2}");
  //RooRealVar P("P","",0.0,5.0);
   RooRealVar P_px("P_px","",0.0,5.0);
   RooRealVar P_py("P_py","",0.0,5.0);
   RooRealVar P_pz("P_pz","",0.0,5.0);
   RooRealVar pi_px("pi_px","",0.0,5.0);
   RooRealVar pi_py("pi_py","",0.0,5.0);
   RooRealVar pi_pz("pi_pz","",0.0,5.0);
   RooRealVar P_E("P_E","",0.0,10.0);
   RooRealVar pi_E("pi_E","",0.0,10.0);   
   //RooRealVar pi_dr("pi_dr","",-30.0,30.0);   
   RooRealVar Distance("Distance","",0.0,40.0);
   RooRealVar Cosalpha("Cosalpha","",-1.0,1.0);
   RooRealVar P_cosTheta("P_cosTheta","",-1.0,1.0);   
   RooRealVar pi_pSVD("pi_pSVD","momentum for pi",0.0,5.0);
   RooRealVar P_pSVD("P_pSVD","momentum for p",0.0,5.0);      
   RooRealVar pi_SVDdEdx("pi_SVDdEdx","",0.0,5.0);
   RooRealVar P_SVDdEdx("P_SVDdEdx","",0.0,5.0);
   RooRealVar pi_p("pi_p","momentum for mu1",0.0,5.0);
   RooRealVar P_p("P_p","momentum for mu2",0.0,5.0);   
   RooRealVar P_nCDCHits("P_nCDCHits","",0.0,500.0); 
   RooRealVar pi_nCDCHits("pi_nCDCHits","",0.0,500.0); 
   RooRealVar P_SVD_nHits("P_SVD_nHits","",0.0,500.0); 
   RooRealVar pi_SVD_nHits("pi_SVD_nHits","",0.0,500.0);
   RooRealVar IsSignal("IsSignal","",0.0,1.0);           
   RooRealVar ChiProb("ChiProb","",0.0,1.0);     
   
   auto variables = new RooArgSet();
   
   variables->add(invm);//
   //variables->add(P);
   variables->add(P_px);
   variables->add(P_py);
   variables->add(P_pz);
   variables->add(pi_px);
   variables->add(pi_py);
   variables->add(pi_pz);
   variables->add(P_E);
   variables->add(pi_E);   
   //variables->add(pi_dr);   
   variables->add(Distance);
   variables->add(Cosalpha);
   variables->add(P_cosTheta);    
   variables->add(pi_pSVD);//
   variables->add(P_pSVD);   //
   variables->add(pi_SVDdEdx);//
   variables->add(P_SVDdEdx);//
   variables->add(pi_p);
   variables->add(P_p);   
   variables->add(ChiProb);
   variables->add(P_nCDCHits);
   variables->add(pi_nCDCHits);
   variables->add(P_SVD_nHits);////
   variables->add(pi_SVD_nHits);////
   variables->add(IsSignal);   //
   //variables->add(Mpipi);
   //variables->add(Mee); 
   
   RooDataSet* data = new RooDataSet("data", "data",*variables);
   
   for(int i=0;i<chain->GetEntries();i++){
        
        chain->GetEntry(i);
        
        //pP_over_pL.setVal(Mom_ratio); //
	invm.setVal(InvM);
	//P.setVal(p); //
	P_px.setVal(p_px);
	P_py.setVal(p_py);
	P_pz.setVal(p_pz);
	pi_px.setVal(Pi_px);
	pi_py.setVal(Pi_py);
	pi_pz.setVal(Pi_pz);
	P_E.setVal(p_E);
	pi_E.setVal(Pi_E);	     
	//pi_dr.setVal(Pi_dr);	     
	Distance.setVal(distance);
	Cosalpha.setVal(cosalpha);
	P_cosTheta.setVal(p_cosTheta);	     
        pi_pSVD.setVal(Pi_pSVD);
	P_pSVD.setVal(p_pSVD);	     
	pi_SVDdEdx.setVal(Pi_SVDdEdx*0.000001);
	P_SVDdEdx.setVal(p_SVDdEdx*0.000001);
	pi_p.setVal(Pi_p);
	P_p.setVal(p_p);	     
	ChiProb.setVal(chiProb);
	P_nCDCHits.setVal(p_nCDCHits);
	pi_nCDCHits.setVal(Pi_nCDCHits); 
	P_SVD_nHits.setVal(p_nSVDHits);
	pi_SVD_nHits.setVal(Pi_nSVDHits);  
	IsSignal.setVal(isSignal);        
	//Mpipi.setVal(m_PiPi); //
	//Mee.setVal(m_ee); //
    
        data->add(*variables); 
        
   }
     RooPlot *myframe = data->plotOn(invm.frame(130));
     
     RooRealVar mean1("mean1", " mean1", 1.116254, 1.1105, 1.122);
     RooRealVar sigma1("sigma1", "#sigma_{1}", 0.004077, 0., 0.01); 
     RooGaussian gauss1("gauss1", "gauss1", invm, mean1, sigma1);
     
     
     RooRealVar r1("r1", "r1", 0.424,0.0,1.0);
     RooRealVar r2("r2", "r2", 0.3262,0.0,1.0);
     RooFormulaVar sigmaL1("sigmaL1", "r1*sigma1",RooArgSet(r1, sigma1));
     RooFormulaVar sigmaR1("sigmaR1", "r2*sigma1",RooArgSet(r2, sigma1));
     RooBifurGauss Bgauss1("Bgauss1", "BGauss1", invm, mean1, sigmaL1,sigmaR1);
     
     
     RooRealVar d("d", "d", -0.0003919,-0.01, 0.01);
     RooRealVar r3("r3", "r3", 0.1767,0.0,1.0);
     RooRealVar r4("r4", "r4", 0.1492,0.0,1.0);
     
     RooFormulaVar mean2("mean2", "mean1+d",RooArgSet(mean1, d));
     RooFormulaVar sigmaL2("sigmaL2", "r3*sigma1",RooArgSet(r3, sigma1));
     RooFormulaVar sigmaR2("sigmaR2", "r4*sigma1",RooArgSet(r4, sigma1));
     RooBifurGauss Bgauss2("Bgauss2", "BGauss2", invm, mean2, sigmaL2,sigmaR2);
     
     
     RooRealVar frac1("frac1", "frac1", 0.3185,0.0,1.0 );
     RooRealVar frac2("frac2", "frac2", 0.8635,0.0,1.0 );
     
     RooAddPdf model1("model1","Bgauss1 + Bgauss2 ",RooArgList(Bgauss1,Bgauss2),RooArgList(frac1));
     
     RooAddPdf signal("signal","model1 + gauss1",RooArgList(model1,gauss1),RooArgList(frac2));
      
     
     //Bkg
     RooRealVar c0("c0","c0",0.094,-2.0,2.0);
     RooRealVar c1("c1","c1",-0.6929,-2.0, 2.0);
     RooChebychev bkg("bkg","bkg",invm,RooArgList(c0,c1));
     
     
     RooRealVar nsig("nsig","Nsig",175259, 1000,1000000);
     RooRealVar nbkg("nbkg","Nbkg",2000,100,1000000);
     RooAddPdf total("total","total pdf",RooArgList(signal,bkg),
		     RooArgList(nsig,nbkg));
     
     total.fitTo(*data,Range(1.1,1.13));
     total.plotOn(myframe);
     
     double chisquare = myframe->chiSquare();
     
     total.paramOn(myframe);
     
     RooHist* hpull = myframe->pullHist() ;
     
     total.plotOn(myframe, Components("signal"), LineColor(kRed),LineStyle(kDashed), LineWidth(2));
     total.plotOn(myframe, Components("bkg"), LineColor(kGreen),LineStyle(kDashed), LineWidth(2));
     total.plotOn(myframe, Components("total"), LineColor(kBlue), LineWidth(2));
     
     myframe->GetXaxis()->SetTitle("#Lambda_{mass} (GeV/c^{2})");
     
     
     myframe->Draw();
     
     
     TLatex l;
     l.SetNDC();
     l.SetTextFont(42);
     l.DrawLatex(0.16,0.8,"#splitline{#font[72]{Belle II }}{#splitline{#scale[0.8]{Early Phase III MC}}{#scale[0.7]{#int}Ldt = 5.0fb^{-1}}}");
     l.SetTextFont(42);   
     
          
     /*RooStats::SPlot* sDataX = new RooStats::SPlot("sData","An SPlot",*data,&total, RooArgList(nsig, nbkg) );
     
     cout <<"Yield of signal is " << nsig.getVal() << "From sWeights it is " << sDataX->GetYieldFromSWeight("nsig") <<endl;
     cout <<"Yield of background is " << nbkg.getVal() <<" From sWeights it is " << sDataX->GetYieldFromSWeight("nbkg")<<endl;
     
     
     for(Int_t i=0; i < 20; i++)
       {
	 cout << "signal Weight   " << sDataX->GetSWeight(i,"nsig") 
	      << " bkg Weight   " << sDataX->GetSWeight(i,"nbkg") 
	      << "  Total Weight   " << sDataX->GetSumOfEventSWeight(i) 
	      << endl;
       }
     
      //RooDataSet * dataw_sig = new RooDataSet(data->GetName(), data->GetTitle(),data, *data->get(),0,"nsig_sw") ;
     RooDataSet *dataw_sig = new RooDataSet("tree", "tree",data,*data->get()) ;
     RooDataSet * dataw_bkg  = new RooDataSet(data->GetName(), data->GetTitle(),data, *data->get(),0,"nbkg_sw") ;
     
     //TFile *froot= new TFile("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/Lambda_splot_MC_Mcut.root","recreate");
     //TFile *froot= new TFile("/home/g_vikas_raj/Downloads/root/SVD/splots_2/test_lambda/Lambda_splot_MC_Mcut_with_svdcuts.root","recreate");
     TFile *froot= new TFile("lambda_ubfit_splot.root","recreate");
     froot->cd();
     RooDataSet::setDefaultStorageType(RooAbsData::Tree);
    ((RooTreeDataStore*)(dataw_sig->store())->tree())->Write(); 
    froot->Close();
     //RooDataSet::setDefaultStorageType(RooAbsData::Tree);
     //((RooTreeDataStore*)(dataw_sig->store())->tree())->SetName("tree");
     //((RooTreeDataStore*)(dataw_sig->store())->tree())->Write();
     froot->Close(); 
         
     
          
     RooPlot* frame_sig_P_p = P_p.frame(100) ;
     frame_sig_P_p->SetTitle("sPlot for the signal (MC)");
     dataw_sig->plotOn(frame_sig_P_p, RooFit::DataError(RooAbsData::SumW2) ) ;
     //frame_sig_P_p->Draw();
          
     RooPlot* frame_bkg_P_p = P_p.frame(100) ;
     frame_bkg_P_p->SetTitle("sPlot for the background (MC)");
     dataw_bkg->plotOn(frame_bkg_P_p, RooFit::DataError(RooAbsData::SumW2) ) ;
     //frame_bkg_P_p->Draw();*/
        
     cout<<"chisquare = "<<chisquare<<endl; 
     cout<<"All done !!!!"<<endl; 
       
     
}