#include <iostream>
#include <fstream>
#include <TMath.h>

using namespace std;
using namespace TMath;
using namespace RooFit;


void WeightedFit(int i, TFile *file, TCanvas *c1){
  TFile *in = new TFile("Example_Migrad.root");
  double rangeEn[2] = {1207,486};
  RooRealVar En("Energy","Energy",rangeEn[i],rangeEn[i]+100);
  RooRealVar w1("w1","w1",0);
  RooDataSet *data = (RooDataSet*)in->Get(Form("final_%d",i));

  //Model Things
  //Generalized Gaussian
  RooRealVar bkg_mean("bkg_mean","bkg_mean",En.getMin()+20,
		      En.getMin(),En.getMin()+40);
  RooRealVar bkg_width("bkg_width","bkg_width",5);
  RooRealVar bkg_fwhm("bkg_fwhm","bkg_fwhm",7,0.5,10);
  RooGenericPdf gen1("gen1","TMath::Erf((bkg_width+Energy-bkg_mean)/(TMath::Sqrt(2)*bkg_fwhm))+TMath::Erf((bkg_width-Energy+bkg_mean)/(TMath::Sqrt(2)*bkg_fwhm))",
		     RooArgSet(En,bkg_mean,bkg_width,bkg_fwhm));
  //Shelf 
  RooRealVar shelf_mean("shelf_mean","shelf_mean",En.getMin()+20,
			En.getMin(),En.getMin()+40);
  RooRealVar shelf_fwhm("shelf_fwhm","shelf_fwhm",20,0.5,100);
  RooGenericPdf shelf("shelf","-TMath::Erf((Energy-shelf_mean)/shelf_fwhm)+1",
		      RooArgSet(En,shelf_mean,shelf_fwhm));
  //Line
  RooRealVar lin_var_1("lin_var_1","lin_var_1",0.5,0,10);
  RooRealVar lin_var_2("lin_var_2","lin_var_2",0.5,0,10);
  RooBernstein bern("bern","bern",En,RooArgList(lin_var_1,lin_var_2));

  //yields
  RooRealVar bkgy1("bkgy1","bkgy1",10,0,1000);
  RooRealVar bkgy2("bkgy2","bkgy2",10,0,1000);
  RooRealVar bkgy3("bkgy3","bkgy3",10,0,1000);
  RooAddPdf model(Form("model_%d",i),Form("model_%d",i),
		  RooArgList(gen1,bern,shelf),
		  RooArgList(bkgy1,bkgy2,bkgy3));

  //Unbinned likelihood
  RooAbsReal *nll = model.createNLL(*data);

  RooMinimizer min(*nll);

  int result_migrad = min.migrad();
  int result_hesse = min.hesse();

  file->cd();
  RooFitResult *res = min.save();
  res->Print();
  res->SetName(Form("Result_%d",i));
  res->SetTitle(Form("Result_%d",i));

  RooWorkspace *w = new RooWorkspace(Form("w%d",i),Form("Workspace_%d",i));
  w->import(model);
  w->import(*res);
  w->Write();

  c1->cd();
  RooPlot *P = En.frame();
  data->plotOn(P,Binning(20));
  model.plotOn(P);
  P->SetTitle(Form("Fit %d",i));
  P->Draw();

  in->Close();
}

void Fits(){
  TFile *out = new TFile("Example_Models.root","RECREATE");
  TCanvas *c1 = new TCanvas("c1","c1",700,500);
  c1->Print("Example_Models.pdf[","pdf");
  WeightedFit(0,out,c1);
  c1->Print("Example_Models.pdf","pdf");
  WeightedFit(1,out,c1);
  c1->Print("Example_Models.pdf","pdf");
  c1->Print("Example_Models.pdf]","pdf");
  c1->Close();
  out->Close();
}

void GenerateData(int i,TFile *file,TCanvas *c1){
  TFile *in = new TFile("Example_Models.root");
  RooWorkspace *w = (RooWorkspace*)in->Get(Form("w%d",i));
  RooRealVar *En = w->var("Energy");
  RooAbsPdf *model = w->pdf(Form("model_%d",i));
  file->cd();
  RooDataSet *data = model->generate(*En,100);
  data->convertToTreeStore();
  TTree *outtree = (TTree*)data->tree();
  outtree->SetName(Form("data_%d",i));
  outtree->Write();
  RooPlot *P = En->frame(Title(Form("Data %d",i)),Bins(20));
  data->plotOn(P);
  model->plotOn(P);
  c1->cd();
  P->Draw();
  in->Close();
}

void Generator(){
  TFile *out = new TFile("Example_Data.root","RECREATE");
  TCanvas *c1 = new TCanvas("c1","c1",700,500);
  c1->Print("Example_Data.pdf[","pdf");
  GenerateData(0,out,c1);
  c1->Print("Example_Data.pdf","pdf");
  GenerateData(1,out,c1);
  c1->Print("Example_Data.pdf","pdf");
  c1->Print("Example_Data.pdf]","pdf");
  c1->Close();
  out->Close();
}

void SimFit(){
  TFile *in = new TFile("Example_Data.root");
  RooRealVar En("Energy","Energy",0,3000);
  double rangeEn[2] = {1207,486};
  RooDataSet *data[2];
  TTree *tree[2];
  for(int i=0; i<2; i++){
    //    En[i] = new RooRealVar(Form("Energy_%d",i),Form("Energy_%d",i),
    //RangeEn[i],RangeEn[i]+100);
    tree[i] = (TTree*)in->Get(Form("data_%d",i));
    data[i] = new RooDataSet(Form("rdata%d",i),Form("rdata%d",i),
			     RooArgSet(En),Import(*tree[i]));
  }
  RooCategory region("regions","regions");
  for(int i=0; i<2; i++){
    region.defineType(Form("region_%d",i));
    En.setRange(Form("fitregion_%d",i),rangeEn[i],rangeEn[i]+100);
    region.setRange(Form("fitregion_%d",i),Form("region_%d",i));
  }
  map<string,RooDataSet*> rmap;
  for(int i=0; i<2; i++){
    rmap[Form("region_%d",i)] = (RooDataSet*)data[i];
  }
  RooDataSet combData("combData","combinedData",En,Index(region),Import(rmap));

  //PDF 0
  RooRealVar bkg_mean_0("bkg_mean_0","bkg_mean_0",rangeEn[0]+20,
		      rangeEn[0],rangeEn[0]+40);
  RooRealVar bkg_width_0("bkg_width_0","bkg_width_0",5);
  RooRealVar bkg_fwhm_0("bkg_fwhm_0","bkg_fwhm_0",7,0.5,10);
  RooGenericPdf gen0("gen0","TMath::Erf((bkg_width_0+Energy-bkg_mean_0)/(TMath::Sqrt(2)*bkg_fwhm_0))+TMath::Erf((bkg_width_0-Energy+bkg_mean_0)/(TMath::Sqrt(2)*bkg_fwhm_0))",
		     RooArgSet(En,bkg_mean_0,bkg_width_0,bkg_fwhm_0));
  //Shelf 
  RooRealVar shelf_mean_0("shelf_mean_0","shelf_mean_0",rangeEn[0]+20,
			rangeEn[0],rangeEn[0]+40);
  RooRealVar shelf_fwhm_0("shelf_fwhm_0","shelf_fwhm_0",20,0.5,100);
  RooGenericPdf shelf_0("shelf_0","-TMath::Erf((Energy-shelf_mean_0)/shelf_fwhm_0)+1",
		      RooArgSet(En,shelf_mean_0,shelf_fwhm_0));
  //Line
  RooRealVar lin_0_1("lin_0_1","lin_0_1",0.5,0,10);
  RooRealVar lin_0_2("lin_0_2","lin_0_2",0.5,0,10);
  RooBernstein bern_0("bern_0","bern_0",En,RooArgList(lin_0_1,lin_0_2));

  //yields
  RooRealVar bkgy1_0("bkgy1_0","bkgy1_0",10,0,1000);
  RooRealVar bkgy2_0("bkgy2_0","bkgy2_0",10,0,1000);
  RooRealVar bkgy3_0("bkgy3_0","bkgy3_0",10,0,1000);
  RooAddPdf model_0("model_0","model_0",
		  RooArgList(gen0,bern_0,shelf_0),
		  RooArgList(bkgy1_0,bkgy2_0,bkgy3_0));
  //PDF 1
  RooRealVar bkg_mean_1("bkg_mean_1","bkg_mean_1",rangeEn[1]+20,
		      rangeEn[1],rangeEn[1]+40);
  RooRealVar bkg_width_1("bkg_width_1","bkg_width_1",5);
  RooRealVar bkg_fwhm_1("bkg_fwhm_1","bkg_fwhm_1",7,0.5,10);
  RooGenericPdf gen1("gen1","TMath::Erf((bkg_width_1+Energy-bkg_mean_1)/(TMath::Sqrt(2)*bkg_fwhm_1))+TMath::Erf((bkg_width_1-Energy+bkg_mean_1)/(TMath::Sqrt(2)*bkg_fwhm_1))",
		     RooArgSet(En,bkg_mean_1,bkg_width_1,bkg_fwhm_1));
  //Shelf 
  RooRealVar shelf_mean_1("shelf_mean_1","shelf_mean_1",rangeEn[1]+20,
			rangeEn[1],rangeEn[1]+40);
  RooRealVar shelf_fwhm_1("shelf_fwhm_1","shelf_fwhm_1",20,0.5,100);
  RooGenericPdf shelf_1("shelf_1","-TMath::Erf((Energy-shelf_mean_1)/shelf_fwhm_1)+1",
		      RooArgSet(En,shelf_mean_1,shelf_fwhm_1));
  //Line
  RooRealVar lin_1_1("lin_1_1","lin_1_1",0.5,0,10);
  RooRealVar lin_1_2("lin_1_2","lin_1_2",0.5,0,10);
  RooBernstein bern_1("bern_1","bern_1",En,RooArgList(lin_1_1,lin_1_2));

  //yields
  RooRealVar bkgy1_1("bkgy1_1","bkgy1_1",10,0,1000);
  RooRealVar bkgy2_1("bkgy2_1","bkgy2_1",10,0,1000);
  RooRealVar bkgy3_1("bkgy3_1","bkgy3_1",10,0,1000);
  RooAddPdf model_1("model_1","model_1",
		  RooArgList(gen1,bern_1,shelf_1),
		  RooArgList(bkgy1_1,bkgy2_1,bkgy3_1));

  RooSimultaneous sim("simPDF","simPDF",region);
  sim.addPdf(model_0,"region_0");
  sim.addPdf(model_1,"region_1");

  RooAbsReal *nll = sim.createNLL(combData);

  RooMinimizer min(*nll);

  int result_migrad = min.migrad();
  int result_hesse = min.hesse();

  RooFitResult *res = min.save();
  res->Print("v");

}