#include "RooGlobalFunc.h"
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooChebychev.h"
#include "RooAddPdf.h"
#include "RooExtendPdf.h"
#include "TCanvas.h"
#include "RooPlot.h"
using namespace RooFit ;


void testfit_weight()
{


  RooRealVar weight("weight","event weight",2.0,-5.0,5.0) ;
  RooRealVar Bmass_weight("Bmass_weight","reconstructed top mass",100,0.0,500) ;
  RooDataSet data_Noweight("data_Noweight","data_Noweight",RooArgSet(Bmass_weight));
  RooDataSet data_weight("data_weight","data_weight",RooArgSet(Bmass_weight, weight), "weight");


  // S e t u p   c o m p o n e n t   p d f s 
  // ---------------------------------------

  // Create Gaussian PDF
  RooRealVar mean("mean","mean of gaussian",100,80.0,120) ;
  RooRealVar sigma1("sigma1","width of gaussian",10.0,5.0,15.0) ;

  //  RooGaussian sig("sig","Signal component 1",x,mean,sigma1) ;  
  RooGaussian sig("sig","Signal component 1",Bmass_weight,mean,sigma1) ;
  
  // Build Chebychev polynomial p.d.f.  
  RooRealVar a0("a0","a0",0.0) ;
  RooChebychev bkg("bkg","Background",Bmass_weight,RooArgList(a0)) ;

  // C o n s t r u c t   m o d e l 
  // -------------------------------------------------------------------

  // Sum the composite signal and background into an extended pdf nsig*sig+nbkg*bkg
  RooRealVar nsig("nsig","number of signal events",3000.0,2000.,14000.0) ;
  RooRealVar nbkg("nbkg","number of background events",5000.0,4000.0,30000.0) ;
  RooAddPdf  model("model","g+a",RooArgList(bkg,sig),RooArgList(nbkg,nsig)) ;

  // C o n s t r u c t   d a t a   s a m p l e   t e s t 
  // -------------------------------------------------------------------
  for(int i= 0;  i<3000; i++)
    {
      weight =4;
      double tmp_sig = gRandom->Gaus(100.0 , 10.0 );
      Bmass_weight = tmp_sig;
      data_Noweight.add(RooArgSet(Bmass_weight));
      data_weight.add(RooArgSet(Bmass_weight, weight), 4);
    }  
  for(int i= 0; i<5000; i++)
    {
      double tmp_bkg = gRandom->Uniform(0.0 , 500.0 );
      weight =4;
      Bmass_weight = tmp_bkg;
      data_Noweight.add(RooArgSet(Bmass_weight));
      data_weight.add(RooArgSet(Bmass_weight, weight), 4);
    }  

  //  F i t   S a m p l e s
  // -------------------------------------------------------------------
  
     // Fit model to data without weight
     // -------------------------------------------------------------------
     cout<<""<<endl;
     cout<<"= = = No weight : = = = "<<endl;
     model.fitTo(data_Noweight) ;

     // Plot No Weight
     TCanvas *Plot_Noweight = new TCanvas();
     RooPlot *RooPlot_Noweight = Bmass_weight.frame();
     RooPlot_Noweight  = Bmass_weight.frame(Name(""), Title(""),Bins(100));
     RooPlot_Noweight->SetTitle("");
     RooPlot_Noweight->GetXaxis()->SetTitle("Mass (MeV/c^{2})");
     RooPlot_Noweight->GetYaxis()->SetTitle("Events / 5 MeV/c^{2}");
     RooPlot_Noweight->GetYaxis()->SetTitleOffset(1.3);
     data_Noweight.plotOn(RooPlot_Noweight ,Name("hist_Noweight"));
     model.plotOn(RooPlot_Noweight, LineColor(kRed),Name("pdf_Noweight"));
     data_Noweight.plotOn(RooPlot_Noweight ,Name("hist_Noweight"));
     RooPlot_Noweight->Draw();
 
     Plot_Noweight->SaveAs("test_NoWeight.C");

  
     // Fit model to data with weight but uncorrected fit
     // -------------------------------------------------------------------
     cout<<""<<endl;
     cout<<"= = = weight uncorrected fit: = = = "<<endl;
     model.fitTo(data_weight, SumW2Error(kFALSE)) ;

     // Plot Weight with Uncorrected fit
     TCanvas *Plot_Uncorrected_w = new TCanvas();
     RooPlot *RooPlot_weight = Bmass_weight.frame();
     RooPlot_weight  = Bmass_weight.frame(Name(""), Title(""),Bins(100));
     RooPlot_weight->SetTitle("");
     RooPlot_weight->GetXaxis()->SetTitle("Mass (MeV/c^{2})");
     RooPlot_weight->GetYaxis()->SetTitle("Events / 5 MeV/c^{2}");
     RooPlot_weight->GetYaxis()->SetTitleOffset(1.3);
     data_weight.plotOn(RooPlot_weight ,Name("hist_weight_NoCorr"));
     model.plotOn(RooPlot_weight, LineColor(kRed),Name("pdf_weight"));
     data_weight.plotOn(RooPlot_weight ,Name("hist_weight_NoCorr"));
     RooPlot_weight->Draw();
 
     Plot_Uncorrected_w->SaveAs("test_Uncorrected_W.C");



     // Fit model to data with weight but corrected fit
     // -------------------------------------------------------------------
     cout<<""<<endl;
     cout<<"= = = weight corrected fit: = = = "<<endl;
     model.fitTo(data_weight, SumW2Error(kTRUE)) ;

     // Plot Weight with corrected fit
     TCanvas *Plot_Corrected_w = new TCanvas();
     RooPlot *RooPlot_weight_Corr = Bmass_weight.frame();
     RooPlot_weight_Corr  = Bmass_weight.frame(Name(""), Title(""),Bins(100));
     RooPlot_weight_Corr->SetTitle("");
     RooPlot_weight_Corr->GetXaxis()->SetTitle("Mass (MeV/c^{2})");
     RooPlot_weight_Corr->GetYaxis()->SetTitle("Events / 5 MeV/c^{2}");
     RooPlot_weight_Corr->GetYaxis()->SetTitleOffset(1.3);
     data_weight.plotOn(RooPlot_weight_Corr ,Name("hist_weight_Corr"), DataError(RooAbsData::SumW2));
     model.plotOn(RooPlot_weight_Corr, LineColor(kRed),Name("pdf_weight_Corr"));
     data_weight.plotOn(RooPlot_weight_Corr ,Name("hist_weight_Corr"), DataError(RooAbsData::SumW2));
     RooPlot_weight_Corr->Draw();
 
     Plot_Corrected_w->SaveAs("test_Corrected_W.C");


}