#include <vector>
#include <iostream>
#include <sstream>
#include <istream>
#include <TFile.h>
#include <TROOT.h>
#include <TH1.h>
#include <TTree.h>
#include <TPaveText.h>
#include <TAxis.h>
#include <TCanvas.h>
#include <TSystem.h>
#include <TStyle.h>
#include <TString.h>
#include <RooRealVar.h>
#include <RooDataSet.h>
#include <RooDataHist.h>
#include <RooGaussian.h>
#include <RooChebychev.h>
#include <RooPolynomial.h>
#include <RooCBShape.h>
#include <RooExponential.h>
#include <RooAddPdf.h>
#include <RooGenericPdf.h>
#include <RooProdPdf.h>
#include <RooEffProd.h>
#include <RooFitResult.h>
#include <RooPlot.h>
#include <RooHist.h>
#include <RooMCStudy.h>
#include <RooRealIntegral.h>
#include "RooHistPdf.h"
#include <TIterator.h>

void fit(string fileName){    

   using namespace RooFit;

   //Open data file
   TFile* inputFile = TFile::Open(fileName.c_str());
   TTree* inputTree;
   RooRealVar M_OS("M_OS","M_OS",500,630,"MeV/c^{2}");
   M_OS.setBins(130);
   TList* keyList = inputFile->GetListOfKeys();
   char treeName[50];
   sprintf (treeName, "%s", keyList->At(0)->GetName()); 
   inputTree = (TTree*)inputFile->Get(treeName);
    
   //Open Ds2eta shape file
   TFile* histfile = TFile::Open("~/Downloads/Ds2eta.root");
   TH1D *hist = (TH1D*) histfile->Get("hetaDs");
    
   RooDataHist *pdfh = new RooDataHist("pdfh", "pdfh", RooArgList(M_OS), hist);

   //Set ranges
   M_OS.setRange("sx",500,544);
   M_OS.setRange("dx",552,630);

   //Import data from TTree and create a RooDataHist
   RooDataSet *data = new RooDataSet("data","data",RooArgList(M_OS),Import(*inputTree),CutRange("sx,dx"));
   RooDataHist *datah = new RooDataHist("datah","datah",RooArgList(M_OS),*data);

   Int_t nEntries = data->numEntries();
   cout << "Number of entries = " << nEntries << endl;

   //The mass plot
   RooPlot *massplot = M_OS.frame();
   datah->plotOn(massplot,Name("datahistogram"));

    
   // **********************THE MASS FIT********************************
   // **********************THE PARAMETERS***************************
    
   RooRealVar numBkg("numBkg","numBkg", 2.0e+05, 1.4321e+05, 2.8321e+05); 
   RooRealVar bkg_p0("bkg_p0","bkg_p0",-0.1,-1,1);
   RooRealVar bkg_p1("bkg_p1","bkg_p1",5e-03,-1,1);
   RooRealVar bkg_p2("bkg_p2","bkg_p2", 0 , -1, 1 );
   RooRealVar bkg_p3("bkg_p3","bkg_p3",0,-1,1);
   RooRealVar bkg_p4("bkg_p4","bkg_p4",0,-1,1);
   RooRealVar bkg_p5("bkg_p5","bkg_p5",-0.0006,-1,1);

   RooRealVar numHist("numHist","numHist", 1700 );

    
   // **********************THE PDFS*********************************
    
   RooHistPdf* histpdf = new RooHistPdf("histpdf","histpdf",M_OS,*pdfh,5);

   RooChebychev* bkg = new RooChebychev("bkg","bkg", M_OS, RooArgList(bkg_p0, bkg_p1, bkg_p2, bkg_p3));

   RooFormulaVar* step = new RooFormulaVar("step", "step", "M_OS < 544 || M_OS > 552", RooArgList(M_OS));
   RooFormulaVar* step2 = new RooFormulaVar("step2", "step2", "M_OS < 544 || M_OS > 552", RooArgList(M_OS));

   RooEffProd* blind_bkg = new RooEffProd("blind_bkg", "blind_bkg", *bkg, *step);
   RooEffProd* blind_eta = new RooEffProd("blind_eta", "blind_eta", *histpdf, *step2);

   RooAddPdf* blind = new RooAddPdf("blind", "blind", RooArgList(*bkg,*histpdf), RooArgList(numBkg,numHist));


   // **********************THE FIT*********************************

   RooFitResult *mass_result;

   mass_result = blind->fitTo(*datah,Extended(),Save(),Range("sx,dx"));
    
   mass_result->Print();

   //Check fit results
   RooArgList list = mass_result->floatParsFinal();
   for (Int_t i = 0; i < list.getSize(); i++)
   {
      RooRealVar* temp = (RooRealVar*)list.at(i);
      if ((temp->getVal() + 3.*temp->getError()) > temp->getMax())
      {
         cout << "Parameter " << temp->GetName() << " exceeds range at 3 sigma!" << endl;
      }
      else if ((temp->getVal() - 3.*temp->getError()) < temp->getMin())
      {
         cout << "Parameter " << temp->GetName() << " exceeds range at 3 sigma!" << endl;
      }
   }

   //We update the mass plot from earlier to show the fit to the data. 
   blind->plotOn(massplot, Components(RooArgSet(*bkg)), LineStyle(5), LineColor(kBlack), Name("fit"));
   blind->plotOn(massplot, Components(RooArgSet(*histpdf)), LineStyle(5), LineColor(kBlack), Name("fit2"));
   blind->plotOn(massplot, LineColor(kBlue), Name("blindtot"));

   massplot->GetXaxis()->SetTitle("#it{m}(#it{#pi}^{ #plus}#it{#pi}^{ #minus}) [MeV/#it{c}^{2}]");
   massplot->GetYaxis()->SetTitle("Candidates / (1 MeV/#it{c}^{2})");
   massplot->SetTitle("");

   massplot->Print("V");
   auto dataHist  = (RooHist*) massplot->getHist("datahistogram"); 
   auto curve1 = (RooCurve*) massplot->getObject(1);  // 1 is index in the list of RooPlot items (see printout from massplot->Print("V")  
   auto curve2 = (RooCurve*) massplot->getObject(2);
   auto hresid1 =  dataHist->makePullHist(*curve1,true);
   auto hresid2 =  dataHist->makePullHist(*curve2,true);

   //RooHist* hresid = massplot->pullHist("datahistogram","blindtot");
   RooPlot* resid = M_OS.frame();
   resid->addPlotable(hresid1,"P"); 
   resid->addPlotable(hresid2,"P"); 
   resid->SetTitle("");
   resid->GetXaxis()->SetTitle("#it{m}(#it{#pi}^{ #plus}#it{#pi}^{ #minus}) [MeV/#it{c}^{2}]");
   gStyle->SetPadLeftMargin(0.1);

   TCanvas *c = new TCanvas("datac","datac",1440,900);
   TPad *pad1 = new TPad("pad1", "The pad 80 of the height",0.0,0.4,1.0,1.0);
   TPad *pad2 = new TPad("pad2", "The pad 20 of the height",0.0,0.0,1.0,0.4);
   pad1->Draw();
   pad2->Draw();
   pad1->cd();
   massplot->Draw();
   pad2->cd();
   resid->Draw();

}


void fitSpectrum(){

   vector<string>  files;
    
   //Set files
   files.push_back("~/Downloads/Spectrum_fit.root");

   for (int i = 0; i < files.size(); i++)
   {
      cout << "Input file = " << files.at(i) << endl;
      fit(files.at(i));
   }
}