#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 the 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_Bi.root");
  TTree* tree1 = (TTree*)file1->Get("output");
  RooDataSet Bi_int("ntuple_Bi_internal","ntuple of internal Bi214", tree1, ntupleVarSet);

  //----Monte Carlo file of category 2 (Tl)---
  TFile *file2 = new TFile("MC_Tl.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_ext.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(200,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(200,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(200,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(200,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); // number of Bi
  RooRealVar nTl("nTl","nTl",400.0,-10000.0,10000.0); // number of Tl
  RooRealVar next("next", "next", 500.0, -10000.0, 10000.0); // number of ext

  // 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));
  
  // Creating category
  RooCategory sample("sample", "sample");
  sample.defineType("b14");
  sample.defineType("E");
  sample.defineType("psnr");
  sample.defineType("udr");

  // Creating combined dataset
  RooDataHist combData("combData", "combined Data", RooArgList(beta14, calenergy, sposr, udotr), Index(sample), Import("b14", hhdatatb1), Import("E", hhdatatb1_2), Import("psnr", hhdatatb1_3), Import("udr", hhdatatb1_4));

  // Creating simultaneous pdf
  RooSimultaneous simpdf("simpdf","simultaneous sample", sample);
  simpdf.addPdf(modelb14,"b14");
  simpdf.addPdf(modelE, "E");
  simpdf.addPdf(modelsposr, "psnr");
  simpdf.addPdf(modeludotr, "udr");

  // Perform fit
  RooFitResult *r = simpdf.fitTo(combData, Extended(), Save());
  

  // Put them into frames for plotting.

  // frame1: Plot all categories in beta14 observable space.
  RooPlot* frame1 = beta14.frame(Bins(200), Title("#beta_{14}"));
  combData.plotOn(frame1, Cut("sample==sample::b14"));
  simpdf.plotOn(frame1, Slice(sample, "b14"), ProjWData(sample, combData), LineWidth(2.0), Name("fit"));
  simpdf.plotOn(frame1, Slice(sample, "b14"), ProjWData(sample, combData), Components("Bibeta14pdf"), LineColor(kRed), LineStyle(8), Name("Bi_int"));
  simpdf.plotOn(frame1, Slice(sample, "b14"), ProjWData(sample, combData), Components("Tlbeta14pdf"), LineColor(kTeal), LineStyle(8), Name("Tl_int"));
  simpdf.plotOn(frame1, Slice(sample, "b14"), ProjWData(sample, combData), Components("extbeta14pdf"), LineColor(kMagenta-5), LineStyle(8), Name("Ext"));
  simpdf.paramOn(frame1, Format("NELU", AutoPrecision(2)));

  // frame2: Plot all categories in calenergy observable space.
  RooPlot* frame2 = calenergy.frame(Bins(50), Title("Calibrated Energy"));
  combData.plotOn(frame2,Cut("sample==sample::E"));
  simpdf.plotOn(frame2,Slice(sample, "E"), ProjWData(sample,combData), LineWidth(2.0));
  simpdf.plotOn(frame2, Slice(sample, "E"), ProjWData(sample, combData), Components("Bienergypdf"), LineColor(kRed), LineStyle(8));
  simpdf.plotOn(frame2, Slice(sample, "E"), ProjWData(sample, combData), Components("Tlenergypdf"), LineColor(kTeal), LineStyle(8));
  simpdf.plotOn(frame2, Slice(sample, "E"), ProjWData(sample, combData), Components("extenergypdf"), LineColor(kMagenta-5), LineStyle(8));

  // frame3: Plot all categories in sposr observable space.
  RooPlot* frame3 = sposr.frame(Bins(100), Title("Radial position"));
  combData.plotOn(frame3,Cut("sample==sample::psnr"));
  simpdf.plotOn(frame3,Slice(sample, "psnr"), ProjWData(sample,combData), LineWidth(2.0));
  simpdf.plotOn(frame3, Slice(sample, "psnr"), ProjWData(sample, combData), Components("Bisposrpdf"), LineColor(kRed), LineStyle(8));
  simpdf.plotOn(frame3, Slice(sample, "psnr"), ProjWData(sample, combData), Components("Tlsposrpdf"), LineColor(kTeal), LineStyle(8));
  simpdf.plotOn(frame3, Slice(sample, "psnr"), ProjWData(sample, combData), Components("extsposrpdf"), LineColor(kMagenta-5), LineStyle(8));
  
}


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

  test4();
  
  theApp.Run(true);


}