#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "RooPlot.h"
#include "RooFitResult.h"
#include "RooGenericPdf.h"
#include "RooConstVar.h"
#include "RooDataHist.h"
#include "RooHistPdf.h"
#include "RooFitResult.h"
#include "RooArgSet.h"
#include "RooAddPdf.h"
#include "RooGaussian.h"
#include "RooCategory.h"
#include "RooSimultaneous.h"
#include "RooChebychev.h"
#include "RooAbsReal.h"
#include "RooMCStudy.h"
#include "RooMinuit.h"
#include "RooNLLVar.h"
#include "RooTrace.h"
#include "RooAddition.h"

#include <TApplication.h>
#include <TH1D.h>
#include <TRandom3.h>
#include <TMinuit.h>
#include <TFile.h>
#include <TTree.h>
#include <TLegend.h>

#include <iostream>
#include <map>

using namespace RooFit;

void test4(){

  // Define lower and upper bounds for variables
  double loE = 4.0;
  double hiE = 6.0;
  double lob14 = -0.2;
  double hib14 = 1.0;
  double lopr = 0.0;
  double hipr = 5500.0;
  double loudr = -1.0;
  double hiudr = 1.0;
  
  // Define variables to match those in ntuples
  RooRealVar calenergy("calenergy","calibrated energy", loE, hiE);
  RooRealVar sposr("sposr", "shifted posr", lopr, hipr);
  RooRealVar beta14("beta14", "beta14 classifier", lob14,hib14);
  RooRealVar udotr("udotr", "recon direction", loudr, hiudr);
  RooArgSet ntupleVarSet(calenergy, sposr, beta14, udotr);

  // Open files and load in ntuples. The ntuples have 4 branches: beta14, calenergy, sposr, udotr
  // --- Monte Carlo file of category 1 (Bi) ----
  TFile *file1 = new TFile("MC_file1.root");
  TTree* tree1 = (TTree*)file1->Get("output");
  RooDataSet Bi_int("ntuple_Bi_internal","ntuple of internal Bi214", tree1, ntupleVarSet);

  //---Monte Carlo file of catergory 2 (Tl)
  TFile *file2 = new TFile("MC_file2.root");
  TTree* tree2 = (TTree*)file2->Get("output");
  RooDataSet Tl_int("ntuple_Tl_internal","ntuple of internal Tl208", tree2, ntupleVarSet);

  //---Monte Carlo file of category 3(ext)
  TFile *file4 = new TFile("MC_file3.root");
  TTree* tree4 = (TTree*)file4->Get("output");
  RooDataSet ext("ntuple_external","ntuple of external", tree4, ntupleVarSet);

  //--- real data file
  TFile *file3 = new TFile("datafile.root");
  TTree* tree3 = (TTree*)file3->Get("output");
  RooDataSet datatb1("datatb1","data of timebin1", tree3, ntupleVarSet);

  // Create histograms from the Monte Carlo ntuples, for the observables for each of the three categories
  // Naming format: category = {Bi, Tl, ext}. obervables = {beta14, calenergy,sposr}. calenergy <-> energy
  TH1 *hBienergy = Bi_int.createHistogram("hBienergy",calenergy, Binning(100, loE, hiE));
  TH1 *hBisposr = Bi_int.createHistogram("hBisposr", sposr, Binning(100,lopr,hipr));
  TH1 *hBibeta14 = Bi_int.createHistogram("hBibeta14", beta14, Binning(24,lob14,hib14));
  TH1 *hBiudotr = Bi_int.createHistogram("hBiudotr", udotr, Binning(20,loudr,hiudr));
  
  TH1 *hTlenergy = Tl_int.createHistogram("hTlenergy",calenergy, Binning(100, loE, hiE));
  TH1 *hTlsposr = Tl_int.createHistogram("hTlsposr", sposr, Binning(100,lopr,hipr));
  TH1 *hTlbeta14 = Tl_int.createHistogram("hTlbeta14", beta14, Binning(24,lob14,hib14));
  TH1 *hTludotr = Tl_int.createHistogram("hTludotr", udotr, Binning(20,loudr,hiudr));
  
  TH1 *hextenergy = ext.createHistogram("hextenergy",calenergy, Binning(100, loE, hiE));
  TH1 *hextsposr = ext.createHistogram("hextsposr", sposr, Binning(100,lopr, hipr));
  TH1 *hextbeta14 = ext.createHistogram("hextbeta14", beta14, Binning(24,lob14,hib14));
  TH1 *hextudotr = ext.createHistogram("hextudotr", udotr, Binning(20,loudr,hiudr));

  // Create histograms of the data files, one histogram for each observable.
  // Basically binning the same data into different observable space
  //---data binned into the beta14 observable
  TH1 *hdatatb1 = datatb1.createHistogram("hdatatb1",beta14,Binning(24,lob14,hib14));
  RooDataHist *hhdatatb1 = new RooDataHist("hhdatatb1", "data for timebin1", beta14, Import(*hdatatb1));

  //---data binned into the calenergy observable
  TH1 *hdatatb1_2 = datatb1.createHistogram("hdatatb1_2",calenergy,Binning(100,loE,hiE));
  RooDataHist *hhdatatb1_2 = new RooDataHist("hhdatatb1_2", "data for timebin1 energy", calenergy, Import(*hdatatb1_2));

  //---data binned into the sposr observable
  TH1 *hdatatb1_3 = datatb1.createHistogram("hdatatb1_3",sposr,Binning(100,lopr,hipr));
  RooDataHist *hhdatatb1_3 = new RooDataHist("hhdatatb1_3", "data for timebin1 sposr", sposr, Import(*hdatatb1_3));

  //---data binned into the udotr observable
  TH1 *hdatatb1_4 = datatb1.createHistogram("hdatatb1_4",udotr,Binning(20,loudr,hiudr));
  RooDataHist *hhdatatb1_4 = new RooDataHist("hhdatatb1_4", "data for timebin1 udotr", udotr, Import(*hdatatb1_4));

  // Converting the histogram of the categories into PDFs
  //----------------------
  RooDataHist Bienergyhist("Bienergyhist","Bi214 energy hist",calenergy,hBienergy);
  RooHistPdf Bienergypdf("Bienergypdf", "Bi214 energy pdf", calenergy, Bienergyhist);

  RooDataHist Tlenergyhist("Tlenergyhist","Tl208 energy hist",calenergy,hTlenergy);
  RooHistPdf Tlenergypdf("Tlenergypdf", "Tl208 energy pdf", calenergy, Tlenergyhist);

  RooDataHist extenergyhist("extenergyhist","external energy hist",calenergy,hextenergy);
  RooHistPdf extenergypdf("extenergypdf", "external energy pdf", calenergy, extenergyhist);
  //-----------------------
  RooDataHist Bibeta14hist("Bibeta14hist", "Bi214 beta14 histogram", beta14, hBibeta14);
  RooHistPdf Bibeta14pdf("Bibeta14pdf", "Bi214 beta14 pdf", beta14, Bibeta14hist);

  RooDataHist Tlbeta14hist("Tlbeta14hist", "Tl208 beta14 histogram", beta14, hTlbeta14);
  RooHistPdf Tlbeta14pdf("Tlbeta14pdf", "Tl208 beta14 pdf", beta14, Tlbeta14hist);

  RooDataHist extbeta14hist("extbeta14hist", "external beta14 histogram", beta14, hextbeta14);
  RooHistPdf extbeta14pdf("extbeta14pdf", "external beta14 pdf", beta14, extbeta14hist);
  //-----------------------
  RooDataHist Bisposrhist("Bisposrhist", "Bi214 sposr histogram", sposr, hBisposr);
  RooHistPdf Bisposrpdf("Bisposrpdf", "Bi214 sposr pdf", sposr, Bisposrhist);

  RooDataHist Tlsposrhist("Tlsposrhist", "Tl208 sposr histogram", sposr, hTlsposr);
  RooHistPdf Tlsposrpdf("Tlsposrpdf", "Tl208 sposr pdf", sposr, Tlsposrhist);

  RooDataHist extsposrhist("extsposrhist", "external sposr histogram", sposr, hextsposr);
  RooHistPdf extsposrpdf("extsposrpdf", "external sposr pdf", sposr, extsposrhist);

  //------

  RooDataHist Biudotrhist("Biudotrhist", "Bi214 udotr histogram", udotr, hBiudotr);
  RooHistPdf Biudotrpdf("Biudotrpdf", "Bi214 udotr pdf", udotr, Biudotrhist);

  RooDataHist Tludotrhist("Tludotrhist", "Tl208 udotr histogram", udotr, hTludotr);
  RooHistPdf Tludotrpdf("Tludotrpdf", "Tl208 udotr pdf", udotr, Tludotrhist);

  RooDataHist extudotrhist("extudotrhist", "external udotr histogram", udotr, hextudotr);
  RooHistPdf extudotrpdf("extudotrpdf", "external udotr pdf", udotr, extudotrhist);

    // Quantities I am fitting for
  RooRealVar nBi("nBi","nBi",400.0,-10000.0,10000.0);
  RooRealVar nTl("nTl","nTl",400.0,-10000.0,10000.0);
  RooRealVar next("next", "next", 500.0, -10000.0, 10000.0);

  // Construct summed PDFs for each of the observables
  RooAddPdf modelb14("modelb14", "modelb14", RooArgList(Bibeta14pdf, Tlbeta14pdf, extbeta14pdf), RooArgList(nBi, nTl, next));
  RooAddPdf modelE("modelE", "modelE", RooArgList(Bienergypdf, Tlenergypdf, extenergypdf), RooArgList(nBi, nTl, next));
  RooAddPdf modelsposr("modelsposr", "modelsposr", RooArgList(Bisposrpdf, Tlsposrpdf, extsposrpdf), RooArgList(nBi, nTl, next));
  RooAddPdf modeludotr("modeludotr", "modeludotr", RooArgList(Biudotrpdf, Tludotrpdf, extudotrpdf), RooArgList(nBi, nTl, next));
  
  // Construct negative log likelihood for each of the observables using corresponding model and dataset
  RooAbsReal *nll1 = modelb14.createNLL(*hhdatatb1, Range(-0.12,0.95));
  RooAbsReal *nll2 = modelE.createNLL(*hhdatatb1_2);
  RooAbsReal *nll3 = modelsposr.createNLL(*hhdatatb1_3);
  RooAbsReal *nll4 = modeludotr.createNLL(*hhdatatb1_4);
  RooAbsReal *nll5 = modelsunct.createNLL(*hhdatatb1_5);
  
  // Construct sum total negative log likelihood and minimize this quantity
  RooAddition nll("nll","nll", RooArgSet(*nll1,*nll2, *nll3, *nll4));

  // Minimize nll
  RooMinuit m(nll);
  m.setVerbose(kTRUE);
  m.migrad();
  m.hesse();
  
  RooFitResult* r = m.save();

  r->Print("v");
  
  Bi_int.Print("V");
  Tl_int.Print("V");
  ext.Print("V");
  
}


int main(int argc, char **argv){
  TApplication theApp("App", &argc, argv);

  test4();
  
  theApp.Run(true);


}