//Code for fitting recoil mass (side-band fit + S+B fit)

//root -l 3.ROOFIT/CombinedModel.C

#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooChebychev.h"
#include "RooPlot.h"
#include "RooWorkspace.h"


using namespace std;
using namespace RooFit;

void  CombinedModel_v2() {
  
  //Import workspace
    
  TFile meas1("SPlusBModel_240GeV_.root");
  TFile meas2("SPlusBModel_365GeV_.root");    
  RooWorkspace* MEAS_240 = (RooWorkspace*)meas1.Get("w");
  RooWorkspace* MEAS_365 = (RooWorkspace*)meas2.Get("w");
  
  //Get original models

  RooAbsPdf * model_240 = MEAS_240->pdf("modelcons");
  RooAbsPdf * model_365 = MEAS_365->pdf("modelcons");
  
  cout << "models = " << model_240 << " " << model_365 << endl;

  //Import data

  RooDataSet* data_240 = (RooDataSet*)MEAS_240->data("data");
  RooDataSet* data_365 = (RooDataSet*)MEAS_365->data("data");

  data_240->Print();
  data_365->Print();
 
  // Generate combined workspace

  RooWorkspace* combWS = new RooWorkspace("combWS");
  combWS->import(*model_240, RenameAllNodes("240GeV"), RenameAllVariables("240GeV"));
  combWS->import(*model_365, RenameAllNodes("365GeV"), RenameAllVariables("365GeV"));

  RooCategory* energy = new RooCategory("energy","energy");
  energy -> defineType("240GeV",0);
  energy -> defineType("365GeV",1);
  combWS ->import(*energy);
  //  combWS -> allCats ();
  
  RooArgSet vars_240 = *((RooArgSet*)data_240->get()->clone("vars"));
  //  vars_240.Print();
   

  RooArgSet vars_365 = *((RooArgSet*)data_365->get()->clone("vars"));                                                                                                                                                       
  //  vars_365.Print();
 
  RooFormulaVar wFunc240("weight","weight","(0.769572809226021)",*MEAS_240->var("m_recoil_T")) ;
  RooFormulaVar wFunc365("weight","weight","(0.0828920124233495)",*MEAS_365->var("m_recoil_T")) ;
   
  RooDataSet* combData = new RooDataSet("combData","combined dataset",vars_240,Index(*energy),Import("240GeV",*data_240));  //WeightVar(w)
  combData->Print();
  RooRealVar * weight = (RooRealVar *)combData->addColumn(wFunc240);
  RooDataSet * wcombData = new RooDataSet(combData->GetName(), combData->GetTitle(), combData, *combData->get(), 0, weight->GetName());
  wcombData->Print();
  cout << wcombData->isWeighted() << endl; 

  RooDataSet* combData1 = new RooDataSet("combData1","combined dataset1",vars_365,Index(*energy),Import("365GeV",*data_365));  //WeightVar(w)
  combData1->Print();
  RooRealVar * weight1 = (RooRealVar *)combData1->addColumn(wFunc365);
  RooDataSet * wcombData1 = new RooDataSet(combData1->GetName(), combData1->GetTitle(), combData1, *combData1->get(), 0, weight1->GetName());
  wcombData1->Print();
  cout << wcombData1->isWeighted() << endl;
 
  wcombData->append(*wcombData1);  

  wcombData->Print();
  cout << wcombData->isWeighted() << endl;  
 

  // Create joint pdf 
  
  combWS->factory("SIMUL::joint(energy,240GeV=modelcons_240GeV,365GeV=modelcons_365GeV)") ;
       
  combWS -> import(*MEAS_240 -> var("mH"));
  combWS -> import(*MEAS_240 -> var("delta_kl"));
  combWS -> factory("EDIT::combModel(joint,mH_240GeV = mH, mH_365GeV = mH, delta_kl_240GeV = delta_kl, delta_kl_365GeV = delta_kl)");
  
  RooAbsPdf* combModel_split = combWS->pdf("joint");
 
  cout << "model joint split ...." << endl;
  combModel_split->Print("t");

  RooAbsPdf* combModel = combWS->pdf("combModel");
   
  cout << "model combined ...." << endl;
  combModel->Print("t");
  //Combined fit
  
  RooFitResult * r =  combModel -> fitTo(*wcombData,Save());  //SumW2Error(kTRUE)
  
  r->Print();
 
  

}