#include <TH1.h>
#include <TF1.h>
#include <TCanvas.h>

#include <RooRealVar.h>
#include <RooGaussian.h>
#include <RooCategory.h>
#include <RooSimultaneous.h>
//#include <RooDataHist.h>
#include <RooDataSet.h>
#include <RooPlot.h>
#include <RooNLLVar.h>
#include <RooMinuit.h>

#include <TTree.h>
#include <TH1.h>
#include <TRandom.h>
#include <TFile.h>

// saves histo h as ttree
void hist2tree(TH1F *h, TString fout, double emin1, double emax1, double emin2, double emax2, TString tr_name="ESpk") {
  if (!h)
	 Fatal("hist2tree","No histo given!");
  
  // Create a new file and a tree within the file
  TFile *file = new TFile(fout, "RECREATE");
  if (!file)
	 Fatal("hist2tree","Can't open file='%s'!",fout.Data());

  TTree *tree = new TTree(tr_name, "Tree of energy spectrum");
  float E1,E2;
  TBranch *bE1 = tree->Branch("E1", &E1, "E1/F");
  TBranch *bE2 = tree->Branch("E2", &E2, "E2/F");
  
  int ne1=0,ne2=0;
  int bmin1=h->GetXaxis()->FindBin(emin1), bmax1=h->GetXaxis()->FindBin(emax1)
	 ,bmin2=h->GetXaxis()->FindBin(emin2), bmax2=h->GetXaxis()->FindBin(emax2);
  // Loop over all bins of the histogram
  for (int b = bmin1; b <= bmax2; ++b) {
	 int N = h->GetBinContent(b);
	 float bc = h->GetBinCenter(b);
	 float bw = h->GetBinWidth(b);
	 if (b>=bmin1 && b<=bmax1)   // first line
		for (int i = 0; i < N; ++i) {  // Generate N values around the bin center
		  E1 = gRandom->Uniform(bc-0.5*bw,bc+0.5*bw); E2=-1;
		  tree->Fill(); ++ne1;
		  //		  Info("rf_example1","E1=%g",E1);
		}
	 if (b>=bmin2 && b<=bmax2)   // second line
		for (int i = 0; i < N; ++i) {  // Generate N values around the bin center
		  E2 = gRandom->Uniform(bc-0.5*bw,bc+0.5*bw); E1=-1;
		  tree->Fill(); ++ne2;
		  //		  Info("rf_example1","E2=%g",E2);
		}
  }
  // Write the tree to the file and close the file
  tree->Write();
  file->Close();
  delete file;

  TString hn=h->GetName();
  Info("hist2tree","Histo='%s' (ne1=%d, ne2=%d) saved as tree in file='%s'",hn.Data(),ne1,ne2,fout.Data());
} // Matches void hist2tree(TH1F *h, TString fout) {

// try to fit two lines simultaneously
void rf_example1() {

  double emin=100,emax=600; // full range of histo

  double emin1=170, emax1=230  // [emin1,emax1] - range to fit 1 line
	 ,emin2=400, emax2=500;     // [emin2,emax2] - range to fit 2 line

  TString ftree="rf_example1_tree.root" // file to save histo as tree
	 ,tr_name="ESpk"; // name of tree

  // generate 3 gaus lines
  TH1F *h= new TH1F("hMain","Main histo",1024,emin,emax);
  TF1 *f = new TF1("f","gaus",emin,emax);
  f->SetParameters(10,200,10);
  h->FillRandom("f",2000);
  f->SetParameters(10,300,20);
  h->FillRandom("f",5000);
  f->SetParameters(10,450,10);
  h->FillRandom("f",1000);
  h->Draw();
  hist2tree(h,ftree,emin1,emax1,emin2,emax2,tr_name); // convert histo -> tree in file ftree
  //  return;

  // now initiate tree
  TFile *file = new TFile(ftree, "READ");
  if (!file)
	 Fatal("hist2tree","Can't open file='%s'!",ftree.Data());
  TTree *tree;
  file->GetObject(tr_name, tree);
  if (!tree)
	 Fatal("hist2tree","Can't find tree='%s' file='%s'!",tr_name.Data(),ftree.Data());
	 
  // simultenous fit of 2 lines (1 & 2)
  RooRealVar E1("E1","energy range of 1 line",emin1,emax1);
  RooRealVar E2("E2","energy range of 2 line",emin2,emax2);

  RooRealVar P1("P1","Position of 1 line",180,220);
  RooRealVar P2("P2","Positions of 2 line",420,470);
  RooRealVar S("S","sigma of lines",0,20) ;
  RooGaussian L1("Line1","Line 1",E1,P1,S) ;
  RooGaussian L2("Line2","Line 2",E2,P2,S) ;

  RooCategory category("category","Category");
  category.defineType("Line1");
  category.defineType("Line2");
  
  RooSimultaneous model("model","Simultaneous fit of 2 lines",category) ;   
  model.addPdf(L1,"Line1");
  model.addPdf(L2,"Line2");

  // build data set using RooDataSet instead of RooDataHist
  RooFormulaVar cut1("cut1", "Cut for E1", "E1>0", E1);
  RooFormulaVar cut2("cut2", "Cut for E2", "E2>0", E2);
  
  RooDataSet dE1("dE1","Line 1 data",tree,RooArgSet(E1),cut1);
  RooDataSet dE2("dE2","Line 2 data",tree,RooArgSet(E2),cut2);

  // Construct combined dataset in (E1,E2,sample)
  RooDataSet combData("combData", "combined data", RooArgSet(E1,E2), 
							 RooFit::Index(category), RooFit::Import("Line1", dE1), 
							 RooFit::Import("Line2", dE2));

  // Plot results
  RooPlot *frame1 = E1.frame(RooFit::Title("RooDataHist: Line 1"))
	 ,*frame2 = E2.frame(RooFit::Title("RooDataHist: Line 2"));
  
  dE1.plotOn(frame1);
  dE2.plotOn(frame2);

  // fit 
  model.fitTo(combData);

  /*
  RooNLLVar nll("nll","nll",model,combData);
  RooMinuit m(nll);
		
  m.setErrorLevel(0.5*sig_err*sig_err);
  m.migrad();
  m.hesse();
  m.minos();
  */

  RooPlot *frame3 = E1.frame(RooFit::Title("Combined data: Line 1"))
	 ,*frame4 = E2.frame(RooFit::Title("Combined data: Line 2"));

  combData.plotOn(frame3, RooFit::Cut("category==category::Line1"));
  combData.plotOn(frame4, RooFit::Cut("category==category::Line2"));
  model.plotOn(frame3,RooFit::Slice(category, "Line1"), RooFit::ProjWData(category,combData));
  model.plotOn(frame4,RooFit::Slice(category, "Line2"), RooFit::ProjWData(category,combData));

  // Plot results

  TCanvas *c = new TCanvas("rf_sample", "rf_sample", 800, 400);
  c->Divide(2,2);
  c->cd(1);
  frame1->Draw();
  c->cd(2);
  frame2->Draw();
  c->cd(3);
  frame3->Draw();
  c->cd(4);
  frame4->Draw();

} // rf_example1