//a dataset is created using the pdf mu*N_sig*S(x)+N_bkg*B(X),
//where the signal is a gaussian and the background is a uniform distribution

//From the fit on the dataset mu and mean(Higgs mass) are estrapolated.
//the mu and Higgs mass distribution is created by repeating for N_toy times

//The proceedings is applied for various alpha values, which increase the number of events
//produced in the dataset 



///////////////////////////////////
//TO RUN THE MACRO
//1. create directory model_2d
//2. run root -b toy_model_2D.C
////////////////////////////////




#include "RooGaussian.h"
#include "RooUniform.h"
#include "RooAbsPdf.h"
#include "TCanvas.h"
#include "RooPlot.h"
#include "TAxis.h"
#include "TH1.h"
#include "TRandom3.h"
#include "RooAddPdf.h"
#include "stdio.h"

using namespace RooFit ;
using namespace std;


void toy_model_2D(){

  const int dim = 10;
  Double_t alpha[dim] = {1,2,5,10,20,30,50,70,90,100};  //scaling factor for events
  
  Double_t N_toy =10000; //number of toy generated
  TString title[dim]={"alpha=1","alpha=2","alpha=5","alpha=10","alpha=20","alpha=30","alpha=50","alpha=70","alpha=90","alpha=100"};
  
  Double_t MU_obs[dim],MU_obs_err[dim],MU_obs_sigma[dim],MU_obs_sigma_err[dim],HIGGS_obs[dim],HIGGS_obs_err[dim],HIGGS_obs_sigma[dim],HIGGS_obs_sigma_err[dim];
  TRandom3 a,b;
  
  Double_t pull,mu_fit, mu_fit_err,pull_Higgs,Higgs_fit_err,Higgs_fit;
  
  Double_t mu_th=1; //theoretical value of mu
  Double_t Higgs_th=125;
 
  gStyle->SetOptStat("MR");  

  
  //Creating histograms
  TH1F *h_mu_pull = new TH1F("h_mu_pull","pulls of mu",50,-3,3);
  TH1F *h_mu_fit = new TH1F("h_mu_fit","Fitted mu",50,0,2.5);
  TH1F *h_mu_err = new TH1F("h_mu_err","Error on fitted mu",25,0,1);
  TH1F *h_Higgs_pull = new TH1F("h_Higgs_pull","pulls of Higgs mass",50,-3,3);
  TH1F *h_Higgs_fit = new TH1F("h_Higgs_fit","Fitted Higgs mass",50,115,130);
  TH1F *h_Higgs_err = new TH1F("h_Higgs_err","Error on fitted Higgs mass",25,0,2);
  TH2F *h_scatter=new TH2F("h_scatter","Scatter plot",50,0,5,50,115,130);
  TCanvas *mu_results = new TCanvas("mu_results","mu_results",1300,1300);
  TCanvas *Higgs_results = new TCanvas("Higgs_results","Higgs_results",1300,1300);

  for(int i=0;i<dim;i++){ //loop alpha
    
    for(int j=0;j<N_toy;j++){

      //Declaring signal variables
      RooRealVar x("x","x",115,130) ;
      RooRealVar mean("mean","Higgs mass",Higgs_th,120,130);
      RooRealVar sigma("sigma","width of gaussian",2.5) ;
      RooRealVar mu("mu","signal strenght",mu_th,0,5);
      
      
      // Build gaussian p.d.f in terms of x,mean and sigma
      RooGaussian sig("sig","Signal shape",x,mean,sigma); 
      RooUniform bkg("bkg","Background shape",x);   
      
      Double_t N_s=(alpha[i])*(a.Poisson(17.35));
      Double_t N_b=(alpha[i])*(b.Poisson(15.78));
      
      
      //Setting Nsig, Nbkg,mu and n_s
      RooRealVar Nsig("Nsig","signal events",N_s);
      RooRealVar Nbkg("Nbkg","bkg events",N_b);
      RooFormulaVar n_s("n_s","@0*@1",RooArgList(Nsig,mu));
      
      
      //Creating model
      RooAddPdf model("model","model",RooArgList(sig,bkg),RooArgList(n_s,Nbkg));
      
      //Creating dataset
      //if number events not specified creates ev=sig+bkg
      RooDataSet *data = model.generate(x);
      
      //Fit pdf to data
      model.fitTo(*data,PrintLevel(-1));
      mu_fit = mu.getVal();
      mu_fit_err=mu.getError();
      pull=(mu_fit - mu_th)/mu_fit_err;
      Higgs_fit = mean.getVal();
      Higgs_fit_err=mean.getError();
      pull_Higgs=(Higgs_fit - Higgs_th)/Higgs_fit_err;
      
      //Filling Histogram
      h_mu_pull->Fill(pull);
      h_mu_fit->Fill(mu_fit);
      h_mu_err->Fill(mu_fit_err);
      h_Higgs_pull->Fill(pull_Higgs);
      h_Higgs_fit->Fill(Higgs_fit);
      h_Higgs_err->Fill(Higgs_fit_err);
      h_scatter->Fill(mu_fit,Higgs_fit);

    }//end for on toy number

    //Creating pdf to fit on MU and HIGGS distribution
    RooRealVar MU("MU","MU",0,2.5);
    RooRealVar MU_mean("MU_mean","mean ",0,2.5) ;
    RooRealVar MU_sigma("MU_sigma","width of gaussian",0,2.5) ;
    RooGaussian MU_sig("MU_sig","Signal shape",MU,MU_mean,MU_sigma); 
    
    RooRealVar HIGGS("HIGGS","HIGGS",115,130);
    RooRealVar HIGGS_mean("HIGGS_mean"," HIGGS mean ",122,127) ;
    RooRealVar HIGGS_sigma("HIGGS_sigma","width of gaussian",0,5) ;
    RooGaussian HIGGS_sig("HIGGS_sig","Signal shape",HIGGS,HIGGS_mean,HIGGS_sigma); 

    //Importing MU and HIGGS distribution to RooDataHist
    //h_mu_fit->Scale(1/N_toy);
    RooDataHist H_MU("H_MU","H_MU",MU,h_mu_fit);
    MU_sig.fitTo(H_MU,PrintLevel(-1),SumW2Error(kTRUE));
    
    MU_obs[i]=MU_mean.getVal();
    MU_obs_err[i]=MU_mean.getError();
    MU_obs_sigma[i]=MU_sigma.getVal();
    MU_obs_sigma_err[i]=MU_sigma.getError();

    //h_Higgs_fit->Scale(1/N_toy);
    RooDataHist H_HIGGS("H_HIGGS","H_HIGGS",HIGGS,h_Higgs_fit);
    HIGGS_sig.fitTo(H_HIGGS,PrintLevel(-1),SumW2Error(kTRUE));
    HIGGS_obs[i]=HIGGS_mean.getVal();
    HIGGS_obs_err[i]=HIGGS_mean.getError();
    HIGGS_obs_sigma[i]=HIGGS_sigma.getVal();
    HIGGS_obs_sigma_err[i]=HIGGS_sigma.getError();
  

    //Drawing mu and Higgs results
    RooPlot* frame = x.frame(Title("Model Gaus+Uniform"));
    data->plotOn(frame);
    model.plotOn(frame,Normalization(1.0,RooAbsReal::RelativeExpected));
    model.plotOn(frame, Components(bkg),LineStyle(kDashed)) ;
      
    RooPlot *frame_mu = MU.frame(Title("MU fitted"));
    H_MU.plotOn(frame_mu);
    MU_sig.plotOn(frame_mu);
   
    mu_results->Divide(2,2);
    mu_results->cd(1);  frame->Draw();
    mu_results->cd(2);  frame_mu->Draw(); 
    mu_results->cd(3);  h_mu_err->Draw();
    mu_results->cd(4);  h_mu_pull->Draw();
    mu_results->Print("model_2d/Mu_results_"+title[i]+".eps");
    mu_results->Clear();
    
    RooPlot *frame_Higgs = HIGGS.frame(Title("HIGGS fitted"));
    H_HIGGS.plotOn(frame_Higgs);
    HIGGS_sig.plotOn(frame_Higgs);
   
    Higgs_results->Divide(2,2);
    Higgs_results->cd(1);  frame->Draw();
    Higgs_results->cd(2);  frame_Higgs->Draw(); 
    Higgs_results->cd(3);  h_Higgs_err->Draw();
    Higgs_results->cd(4);  h_Higgs_pull->Draw();
    Higgs_results->Print("model_2d/Higgs_results_"+title[i]+".eps");
    Higgs_results->Clear();

    gStyle->SetPalette(1);
    gStyle->SetOptStat("");
    TCanvas *scat=new TCanvas("scat","scat");
    h_scatter->Draw("colz");
    scat->Print("model_2d/scatter_plot_"+title[i]+".eps");

    //Resetting Histogram
    h_mu_pull->Reset();
    h_mu_fit->Reset();
    h_mu_err->Reset();
    h_Higgs_pull->Reset();
    h_Higgs_fit->Reset();
    h_Higgs_err->Reset();
    h_scatter->Reset();
  }//end loop alpha

  //printing fitted values  
  for(int i=0;i<dim;i++){
    cout<<""<<endl;
    cout<<"alpha: "<<alpha[i]<<endl;
    cout<<"                   mu fitted: "<<MU_obs[i]<<" +/- "<<MU_obs_err[i]<<endl;     
    cout<<"                sigma fitted: "<<MU_obs_sigma[i]<<" +/- "<<MU_obs_sigma_err[i]<<endl;       
    cout<<"           Higgs mass fitted: "<<HIGGS_obs[i]<<" +/- "<<HIGGS_obs_err[i]<<endl; 
    cout<<"                sigma fitted: "<<HIGGS_obs_sigma[i]<<" +/- "<<HIGGS_obs_sigma_err[i]<<endl;
   }



} //end void