//Make splots from mPip fit


using namespace RooFit; 
using namespace RooStats; 

void Lambda_gaus()
{
   
   //SetBelle2Style() ;
   
   TChain *chain = new TChain("Lambda");
   
   //chain->Add("dEdxcalibration_MC.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/mixed_5fb_Mcut.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/charged_5fb_Mcut.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/uubar_5fb_Mcut.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/ddbar_5fb_Mcut.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/ssbar_5fb_Mcut.root");
  chain->Add("/home/g_vikas_raj/Downloads/root/SVD/splots_2/mcut/ccbar_5fb_Mcut.root");

   
   
   
   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 ;
   
   chain->SetBranchAddress("M",& M);
   chain->SetBranchAddress("InvM",& InvM);
   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);
   chain->SetBranchAddress("p_pSVD",& p_pSVD);   
   chain->SetBranchAddress("Pi_SVDdEdx",& Pi_SVDdEdx);
   chain->SetBranchAddress("p_SVDdEdx",& p_SVDdEdx);
   chain->SetBranchAddress("Pi_p",& Pi_p);
   chain->SetBranchAddress("p_p",& p_p);   
   chain->SetBranchAddress("chiProb",& chiProb);
   chain->SetBranchAddress("isSignal",& isSignal);
   chain->SetBranchAddress("p_nCDCHits",& p_nCDCHits);
   chain->SetBranchAddress("Pi_nCDCHits",& Pi_nCDCHits);
   chain->SetBranchAddress("p_nSVDHits",& p_nSVDHits);
   chain->SetBranchAddress("Pi_nSVDHits",& Pi_nSVDHits);
   
   
   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);     
   
   RooRealVar Mpipi("Mpipi","M_{#pi^{+}#pi^{-}}",0.0,1.13,"GeV/c^{2}");  
   RooRealVar Mee("Mee","M_{e^{+}e^{-}}",0.0,1.13,"GeV/c^{2}");
   
   auto variables = new RooArgSet();
   
   variables->add(pP_over_pL);
   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);
   
   TParticle *Pi =  new TParticle();
   Pi->SetPdgCode(211);
   
   TParticle *el =  new TParticle();
   el->SetPdgCode(11);
   
   TH1D *h1 = new TH1D("h1","",100,1.1,1.13);
   TH1D *h2 = new TH1D("h2","",100,0.0,4.0);
   TH1D *h3 = new TH1D("h3","",100,0.0,2.0);
   
   for(int i=0;i<chain->GetEntries()*0.01;i++) {
   
     chain->GetEntry(i);
     
     if(InvM>1.10 && InvM<1.13 && distance>1.0 && p_nCDCHits>0.0 && /*Pi_nCDCHits>0.0 &&*/ p_nSVDHits > 0.0 && Pi_nSVDHits > 0.0)     
      {
	 
	 double Mom_ratio = p_p / p ;
	 
         double Pi_2_E = sqrt((p_px*p_px)+(p_py*p_py)+(p_pz*p_pz)+Pi->GetMass()*Pi->GetMass());
         double m_PiPi = sqrt((Pi_2_E+Pi_E)*(Pi_2_E+Pi_E)-(p_px+Pi_px)*(p_px+Pi_px)-(p_py+Pi_py)*(p_py+Pi_py)-(p_pz+Pi_pz)*(p_pz+Pi_pz));
	 
         double e_1_E = sqrt((p_px*p_px)+(p_py*p_py)+(p_pz*p_pz)+el->GetMass()*el->GetMass());
         double e_2_E = sqrt((Pi_px*Pi_px)+(Pi_py*Pi_py)+(Pi_pz*Pi_pz)+el->GetMass()*el->GetMass());
         double m_ee = sqrt((e_1_E+e_2_E)*(e_1_E+e_2_E)-(p_px+Pi_px)*(p_px+Pi_px)-(p_py+Pi_py)*(p_py+Pi_py)-(p_pz+Pi_pz)*(p_pz+Pi_pz));
	 
	 
         if((m_PiPi < 0.488 || m_PiPi > 0.508) && m_ee > 0.050)
           {  
	     
	     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); //1.116254
     RooRealVar sigma1("sigma1", "#sigma_{1}", 0.003077, 0., 0.01); //0.003077
     RooGaussian gauss1("gauss1", "gauss1", invm, mean1, sigma1);   
     
     
     RooRealVar r1("r1", "r1", 0.424,0.0,1.0); //0.424 //0.424 looks good?
     RooRealVar r2("r2", "r2", 0.23262,0.0,1.0); //0.3262 //23262 looks good?
     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 sigmaL1("sigmaL1", "sigmaL1",0.000438, 0., 0.01); //0.000438
     RooRealVar sigmaR1("sigmaR1", "sigmaR1",0.000408, 0., 0.01); //0.000408
     RooBifurGauss Bgauss1("Bgauss1", "BGauss1", invm, mean1, sigmaL1,sigmaR1);*/     
     
     RooRealVar d("d", "d", -0.0003919,-0.01, 0.01); //-0.0003919
     RooRealVar r3("r3", "r3", 0.3767,0.0,1.0); //0.1767 //0.3767 looks good?
     RooRealVar r4("r4", "r4", 0.2492,0.0,1.0); //0.1492 //2492 looks good?
     
     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 mean2("mean2", "mean2",1.1158621, 1.1105, 1.122); //1.1158621
     RooRealVar sigmaL2("sigmaL2", "sigmaL2",0.000544, 0., 0.01); //0.000544
     RooRealVar sigmaR2("sigmaR2", "sigmaR2",0.000460, 0., 0.01); //0.000460
     RooBifurGauss Bgauss2("Bgauss2", "BGauss2", invm, mean2, sigmaL2,sigmaR2);*/
     
     
     RooRealVar frac1("frac1", "frac1", 0.3815,0.0,1.0 ); //0.3815
     RooRealVar frac2("frac2", "frac2", 0.8635,0.0,1.0 ); //0.8635
     
     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.1,-2.0,2.0); //0.094
     RooRealVar c1("c1","c1",-0.7,-2.0, 2.0); //-0.6929
     RooChebychev bkg("bkg","bkg",invm,RooArgList(c0,c1));
     
     
     RooRealVar nsig("nsig","Nsig",175259, 1000,1000000); //175259
     RooRealVar nbkg("nbkg","Nbkg",2000,100,1000000);   //2000
     RooAddPdf total("total","total pdf",RooArgList(signal,bkg), RooArgList(nsig,nbkg));
     //RooAddPdf total("total","total pdf",RooArgList(model1,bkg), RooArgList(nsig,nbkg));
     
     //gauss1.fitTo(*data,Range(1.1,1.13));
     //signal.fitTo(*data,Range(1.1,1.13));
     //model1.fitTo(*data,Range(1.1,1.13));
          
     total.fitTo(*data,Range(1.1,1.13));
     
     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));
     
     //signal.plotOn(myframe);
     //model1.plotOn(myframe, LineColor(kRed),LineStyle(kDashed), LineWidth(2));     
     //gauss1.plotOn(myframe, LineColor(kGreen),LineStyle(kDashed), LineWidth(2));
          
     double chisquare = myframe->chiSquare();
     
     //gauss1.paramOn(myframe);
     //signal.paramOn(myframe);
     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 = 3.1fb^{-1}}}");
     l.SetTextFont(42);   
      
     //myframe->saveOn("x1.png");     
     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");
     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;
}