#include "TRandom3.h"
#include "TH2.h"
#include "TH1.h"
#include "TCanvas.h"
#include "TMath.h"
#include "TTree.h"

//ROOFIT
#include "RooRealVar.h"
#include "RooHistPdf.h"
#include "RooDataHist.h"
#include "RooAddPdf.h"
#include "RooPlot.h"
#include "RooNDKeysPdf.h"
#include "RooDataSet.h"

using namespace RooFit;
using namespace std;

int main(int argc, char **argv){

  //Generate pdf histos and Data to fit
  TRandom3 *rgen = new TRandom3();
  TH2F *hGaus0 = new TH2F("hGaus0","hGaus0",50,-10,10,50,-10,10);
  TH2F *hGaus1 = new TH2F("hGaus1","hGaus1",50,-10,10,50,-10,10);
  TH2F *hData = new TH2F("hData","hData",50,-10,10,50,-10,10);

  //Fill histograms for pdf
  for(int ipoint=0; ipoint<100000; ipoint++)
    {
      hGaus0->Fill( rgen->Gaus(0,2), rgen->Gaus(0,2) );
      hGaus1->Fill( rgen->Gaus(1,1), rgen->Gaus(1,1) );
    }
  
  //Fill histogram/tree that will be fitted
  Double_t X,Y;
  TTree *tree_data = new TTree("tree","tree");
  tree_data->Branch("X",&X,"X/D");
  tree_data->Branch("Y",&Y,"Y/D");
  for(int ipoint=0; ipoint<10000; ipoint++)
    {
      X=rgen->Gaus(0,2);
      Y=rgen->Gaus(0,2);
      hData->Fill( X,Y );
      tree_data->Fill();
    }
  for(int ipoint=0; ipoint<30000; ipoint++)
    {
      X=rgen->Gaus(1,1);
      Y=rgen->Gaus(1,1);
      hData->Fill( X,Y );
      tree_data->Fill();
    }


  //Declare observables
  RooRealVar *_roo_X = new RooRealVar("X","X",-10,10);
  RooRealVar *_roo_Y = new RooRealVar("Y","Y",-10,10);

  //Build PDF
  RooDataHist *_roo_datahist0 = NULL;
  RooDataHist *_roo_datahist1 = NULL;
  RooHistPdf *_roo_histpdf0 = NULL;
  RooHistPdf *_roo_histpdf1 = NULL;
  RooAddPdf *_roo_model = NULL;
  _roo_datahist0 = new RooDataHist( "_roo_datahist0","_roo_datahist0", RooArgSet(*_roo_X,*_roo_Y), Import(*hGaus0) );
  _roo_datahist0->Print();
  _roo_histpdf0 = new RooHistPdf( "_roo_histpdf0", "_roo_histpdf0", RooArgSet(*_roo_X,*_roo_Y),*_roo_datahist0);
  _roo_histpdf0->Print();
  _roo_datahist1 = new RooDataHist( "_roo_datahist1","_roo_datahist1", RooArgSet(*_roo_X,*_roo_Y), Import(*hGaus1) );
  _roo_datahist1->Print();
  _roo_histpdf1 = new RooHistPdf( "_roo_histpdf1","_roo_histpdf1", RooArgSet(*_roo_X,*_roo_Y),*_roo_datahist1);
  _roo_histpdf1->Print();

  //Build composite pdf
  RooRealVar *n0 = new RooRealVar("n0","n0_title",hData->GetEntries()*0.5,0,hData->GetEntries()) ;
  RooRealVar *n1 = new RooRealVar("n1","n1_title",hData->GetEntries()*0.5,0,hData->GetEntries()) ;
  _roo_model = new RooAddPdf("model","fit to data",RooArgList(*_roo_histpdf0,*_roo_histpdf1),RooArgList(*n0,*n1)) ;

  //Build Data to fit
   RooDataHist *_roo_datahist_data = new RooDataHist( "_roo_datahist_data","_roo_datahist_data", RooArgSet(*_roo_X,*_roo_Y), Import(*hData) );
   _roo_datahist_data->Print();
   RooDataSet *_roo_dataset_data = new RooDataSet( "_roo_dataset_data", "roo_dataset_data", RooArgSet(*_roo_X,*_roo_Y), Import(*tree_data) );
   _roo_dataset_data->Print();
   RooDataHist* _roo_datahist_binnedclone = _roo_dataset_data->binnedClone() ;
   _roo_datahist_binnedclone->Print();
   
   //FIT
   //_roo_model->fitTo( *_roo_datahist_data, Extended(kTRUE), Verbose(kTRUE), PrintLevel(3));
   _roo_model->fitTo( *_roo_dataset_data, Extended(kTRUE), Verbose(kTRUE), PrintLevel(3));
   //_roo_model->fitTo( *_roo_datahist_binnedclone, Extended(kTRUE), Verbose(kTRUE), PrintLevel(3));

  //Plot
  RooPlot *_roo_frameX =_roo_X->frame(); 
  RooPlot *_roo_frameY =_roo_Y->frame();
  _roo_dataset_data->plotOn( _roo_frameX );
  _roo_model->plotOn( _roo_frameX, LineColor(kMagenta), Normalization(1.0,RooAbsReal::RelativeExpected));
  _roo_model->plotOn( _roo_frameX, Components( *_roo_histpdf0 ), LineColor(kBlue), Normalization(1.0,RooAbsReal::RelativeExpected));
  _roo_model->plotOn( _roo_frameX, Components( *_roo_histpdf1 ), LineColor(kRed),  Normalization(1.0,RooAbsReal::RelativeExpected));
  _roo_dataset_data->plotOn( _roo_frameY );
  _roo_model->plotOn( _roo_frameY, LineColor(kMagenta), Normalization(1.0,RooAbsReal::RelativeExpected));
  _roo_model->plotOn( _roo_frameY, Components( *_roo_histpdf0 ), LineColor(kBlue), Normalization(1.0,RooAbsReal::RelativeExpected));
  _roo_model->plotOn( _roo_frameY, Components( *_roo_histpdf1 ), LineColor(kRed),  Normalization(1.0,RooAbsReal::RelativeExpected));
  
  TCanvas *c = new TCanvas();
  c->Divide(2,1);
  c->cd(1);
  _roo_frameX->Draw("");
  c->cd(2);
  _roo_frameY->Draw("");
  c->SaveAs("roofit_test2_c.eps");

  
  return 0;

}