#include <iostream>
#include <cstdlib>
#include <math.h>
#include <string.h>

#include "TROOT.h"
#include "TPaveText.h"
#include "TCanvas.h"
#include "TH1F.h"
#include "TAxis.h"
#include "TTree.h"
#include "TCut.h"
#include "TString.h"
#include "TSystem.h"
#include "TStyle.h"
#include "TFile.h"
#include "TLegend.h"
#include "TLatex.h"
#include "TObject.h"
#include "TVirtualPad.h"
#include "TMath.h"
#include "RooArgList.h"
#include "RooMinuit.h"
#include "RooGaussian.h"
#include "RooRelBreitWigner.cc"
#include "RooRelBreitWigner.hh"
#include "RooBreitWigner.h"
#include "RooChebychev.h"
#include "RooNumConvPdf.h"
#include "RooAddPdf.h"
#include "RooPlot.h"
#include "RooHist.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooFitResult.h"
#include "RooExtendPdf.h"
#include "RooRealVar.h"
#include "Constants.hh"

using namespace RooFit;
using namespace std;

void FitDsPiMC12_sPlot(TString sPlotfile = "", TString MCfile = "", TString outputfile = ""){
  
  //Setup LHCb style for plots
  gROOT->SetStyle("Plain");
  gROOT->ProcessLine(".x lhcbStyle.C");
  TPaveText *lhcbNameR = new TPaveText(0.70 - gStyle->GetPadRightMargin(),
                                       0.75 - gStyle->GetPadTopMargin(),
                                       0.95 - gStyle->GetPadRightMargin(),
                                       0.85 - gStyle->GetPadTopMargin(),"BRNDC");
  lhcbNameR->AddText("LHCb");
  lhcbNameR->SetFillColor(0);
  lhcbNameR->SetTextAlign(12);
  lhcbNameR->SetBorderSize(0);

  TPaveText *lhcbNameL = new TPaveText(gStyle->GetPadLeftMargin() + 0.05,
                                       0.87 - gStyle->GetPadTopMargin(),
                                       gStyle->GetPadLeftMargin() + 0.20,
                                       0.95 - gStyle->GetPadTopMargin(),"BRNDC");
  lhcbNameL->AddText("LHCb");
  lhcbNameL->SetFillColor(0);
  lhcbNameL->SetTextAlign(12);
  lhcbNameL->SetBorderSize(0);

  //RooMsgService::instance().setGlobalKillBelow(RooFit::WARNING);

  //Define useful RooRealVars  
  RooRealVar DsPiMass("DsPiMass","M(D^{*}#pi_{**})-M(D^{*})+M(D^{*}_{PDG})",min_DsPi,max_DsPi,"GeV/c^{2}");
  
  //Retrieve weigthed m(D*pi) distribution   
  cout << "Retrieving weighted datasets...\n";

  TFile *f_sPlotfile = new TFile(sPlotfile);
  RooWorkspace *w_sPlotLogIP = (RooWorkspace*)f_sPlotfile->Get("w_sPlotLogIP");
  RooAbsData *data_wSGN = w_sPlotLogIP->data("data_wSGN");
  RooAbsData *data_wBKG = w_sPlotLogIP->data("data_wBKG");

  //Compare weighted distributions
  cout << "Comparing weighted distributions...\n";

  RooPlot *frame_comp = DsPiMass.frame(Bins(bin_DsPi),Name("frame"));
  data_wSGN->plotOn(frame_comp,LineColor(kBlack),MarkerColor(kBlack),RooFit::Name("wSGN"));
  data_wBKG->plotOn(frame_comp,LineColor(kRed),MarkerColor(kRed),RooFit::Name("wBKG"));
  
  TCanvas *c_comp = new TCanvas("c_comp","c_comp",950,900);
  frame_comp->Draw();
  TObject *wSGN = frame_comp->findObject("wSGN");
  TObject *wBKG = frame_comp->findObject("wBKG");
  lhcbNameR->Draw("SAME");
  TLegend *leg_comp = new TLegend(0.8,0.8,1.0,1.0);
  leg_comp->AddEntry(wSGN,"Signal","L");
  leg_comp->AddEntry(wBKG,"Combinatorial","L");
  leg_comp->Draw("SAME");
  lhcbNameR->Draw("SAME");

  //-------------------------------------------------------------Fit invariant mass-------------------------------------------------------------
  
  cout << "Fit to invariant mass...\n";
  
  //Get some MC input
  TFile *f_MC = new TFile(MCfile,"READ");

  //Non resonant
  cout << "Non-resonant...\n";

  RooRealVar *c0_ = (RooRealVar*)f_MC->Get("c0_");
  RooRealVar *c1_ = (RooRealVar*)f_MC->Get("c1_");
  RooRealVar *c2_ = (RooRealVar*)f_MC->Get("c2_");
  RooRealVar *c3_ = (RooRealVar*)f_MC->Get("c3_");

  RooRealVar c0("c0","c0",c0_->getVal());
  RooRealVar c1("c1","c1",c1_->getVal());
  RooRealVar c2("c2","c2",c2_->getVal());
  RooRealVar c3("c3","c3",c3_->getVal());
  RooChebychev modelNRpdf("modelNRpdf","modelNRpdf",DsPiMass,RooArgList(c0,c1,c2,c3));
  RooRealVar NNR("NNR","NNR",1000,0,1000000);

  //D1(2420)^0 + resolution
  cout << "D1(2420)^0 + resolution...\n";

  RooRealVar *s1D1_ = (RooRealVar*)f_MC->Get("s1D1_");
  RooRealVar *s2D1_ = (RooRealVar*)f_MC->Get("s2D1_");
  RooRealVar *mD1_ = (RooRealVar*)f_MC->Get("mD1_");
  RooRealVar *gD1_ = (RooRealVar*)f_MC->Get("gD1_");
  RooRealVar *fG1D1_ = (RooRealVar*)f_MC->Get("fG1D1_");

  RooRealVar mD1("mD1","mD1",mD1_->getVal()); 
  RooRealVar gD1("gD1","gD1",gD1_->getVal());
  RooRelBreitWigner bwD1("bwD1","bwD1",DsPiMass,mD1,gD1,r,DsMass,PionMass,spin2);

  RooRealVar s1D1("s1D1","s1D1",s1D1_->getVal());
  RooRealVar s2D1("s2D1","s2D1",s2D1_->getVal());
  RooGaussian G1D1("G1D1","G1D1",DsPiMass,mg,s1D1);
  RooGaussian G2D1("G2D1","G2D1",DsPiMass,mg,s2D1);
  RooRealVar fG1D1("fG1D1","fG1D1",fG1D1_->getVal());
  RooAddPdf resD1("resD1","resD1",RooArgList(G1D1,G2D1),fG1D1);

  RooNumConvPdf modelD1pdf("modelD1pdf","modelD1pdf",DsPiMass,bwD1,G1D1);
  RooRealVar wsD1res("wsD1res","wsD1res",(s1D1.getVal())*(fG1D1.getVal()) + (s2D1.getVal())*(1 - fG1D1.getVal()));
  modelD1pdf.setConvolutionWindow(mg,s1D1,5);

  RooRealVar ND1("ND1","ND1",1000,0,1000000);

  //D2*(2460)^0 + resolution
  cout << "D2*(2460)^0 + resolution...\n";

  RooRealVar *s1D2s_ = (RooRealVar*)f_MC->Get("s1D2s_");
  RooRealVar *s2D2s_ = (RooRealVar*)f_MC->Get("s2D2s_");
  RooRealVar *mD2s_ = (RooRealVar*)f_MC->Get("mD2s_");
  RooRealVar *gD2s_ = (RooRealVar*)f_MC->Get("gD2s_");
  RooRealVar *fG1D2s_ = (RooRealVar*)f_MC->Get("fG1D2s_");

  RooRealVar mD2s("mD2s","mD2s",mD2s_->getVal()); //from PDG
  RooRealVar gD2s("gD2s","gD2s",gD2s_->getVal()); //from PDG
  RooRelBreitWigner bwD2s("bwD2s","bwD2s",DsPiMass,mD2s,gD2s,r,DsMass,PionMass,spin2);

  RooRealVar s1D2s("s1D2s","s1D2s",s1D2s_->getVal());
  RooRealVar s2D2s("s2D2s","s2D2s",s2D2s_->getVal());
  RooGaussian G1D2s("G1D2s","G1D2s",DsPiMass,mg,s1D2s);
  RooGaussian G2D2s("G2D2s","G2D2s",DsPiMass,mg,s2D2s);
  RooRealVar fG1D2s("fG1D2s","fG1D2s",fG1D2s_->getVal());
  RooAddPdf resD2s("resD2s","resD2s",RooArgList(G1D2s,G2D2s),fG1D2s);

  RooNumConvPdf modelD2spdf("modelD2spdf","modelD2spdf",DsPiMass,bwD2s,resD2s);
  RooRealVar wsD2sres("wsD2sres","wsD2sres",(s1D2s.getVal())*(fG1D2s.getVal()) + (s2D2s.getVal())*(1 - fG1D2s.getVal()));
  modelD2spdf.setConvolutionWindow(mg,wsD2sres,5);

  RooRealVar ND2s("ND2s","ND2s",1000,0,1000000);

  //D1(2430)^0 + resolution
  cout << "D1(2430)^0 + resolution...\n";

  RooRealVar *s1D1br_ = (RooRealVar*)f_MC->Get("s1D1br_");
  RooRealVar *s2D1br_ = (RooRealVar*)f_MC->Get("s2D1br_");
  RooRealVar *mD1br_ = (RooRealVar*)f_MC->Get("mD1br_");
  RooRealVar *gD1br_ = (RooRealVar*)f_MC->Get("gD1br_");
  RooRealVar *fG1D1br_ = (RooRealVar*)f_MC->Get("fG1D1br_");

  RooRealVar mD1br("mD1br","mD1br",mD1br_->getVal()); //from PDG
  RooRealVar gD1br("gD1br","gD1br",gD1br_->getVal()); //from PDG
  RooRelBreitWigner bwD1br("bwD1br","bwD1br",DsPiMass,mD1br,gD1br,r,DsMass,PionMass,spin0);

  RooRealVar s1D1br("s1D1br","s1D1br",s1D1br_->getVal());
  RooRealVar s2D1br("s2D1br","s2D1br",s2D1br_->getVal());
  RooGaussian G1D1br("G1D1br","G1D1br",DsPiMass,mg,s1D1br);
  RooGaussian G2D1br("G2D1br","G2D1br",DsPiMass,mg,s2D1br);
  RooRealVar fG1D1br("fG1D1br","fG1D1br",fG1D1br_->getVal());
  RooAddPdf resD1br("resD1br","resD1br",RooArgList(G1D1br,G2D1br),fG1D1br);

  RooNumConvPdf modelD1brpdf("modelD1brpdf","modelD1brpdf",DsPiMass,bwD1br,resD1br);
  RooRealVar wsD1brres("wsD1brres","wsD1brres",(s1D1br.getVal())*(fG1D1br.getVal()) + (s2D1br.getVal())*(1 - fG1D1br.getVal()));
  modelD1brpdf.setConvolutionWindow(mg,wsD1brres,5);

  RooRealVar ND1br("ND1br","ND1br",1000,0,1000000);
  
  //Total model
  cout << "Total pdf...\n";

  RooAddPdf modelDsPiMass("modelDsPiMass","modelDsPiMass"
			  ,RooArgList(modelNRpdf,modelD1pdf,modelD2spdf,modelD1brpdf)
			  ,RooArgList(NNR,ND1,ND2s,ND1br));
  
  modelDsPiMass.printCompactTree();

  //Unbinned maximum likelihood fit
  cout << "Fit to invariant mass...\n";

  RooFitResult *r_DsPiMass = modelDsPiMass.fitTo(*data_wSGN,Range(min_DsPi_fit,max_DsPi_fit),
						 Extended(kTRUE),
						 Save(),
						 SumW2Error(kTRUE),
						 //PrintEvalErrors(-1),
						 //Warnings(kFALSE),
						 Timer(kTRUE));
  
  //Plot
  cout << "Plot results...\n";

  RooPlot *frame = DsPiMass.frame(Bins(bin_DsPi),Name("frame"));
  data_wSGN->plotOn(frame);
  modelDsPiMass.plotOn(frame,RooFit::Name("modelDsPiMass"),LineColor(kBlue));
  RooHist *pull = frame->pullHist(0,0,true);
  RooPlot *frame_pull = DsPiMass.frame(Bins(bin_DsPi),Name("frame_pull"));
  modelDsPiMass.plotOn(frame,Components(modelNRpdf),RooFit::Name("modelNRpdf"),LineColor(kRed),LineStyle(kDashed));
  modelDsPiMass.plotOn(frame,Components(modelD1pdf),RooFit::Name("modelD1pdf"),LineColor(kBlack),LineStyle(kDashed));
  modelDsPiMass.plotOn(frame,Components(modelD2spdf),RooFit::Name("modelD2spdf"),LineColor(kMagenta),LineStyle(kDashed));
  modelDsPiMass.plotOn(frame,Components(modelD1brpdf),RooFit::Name("modelD1brpdf"),LineColor(kPink),LineStyle(kDashed));

  frame_pull->addPlotable(pull,"P");
  frame_pull->SetAxisRange(-5,5,"Y");
  frame_pull->SetNdivisions(12,"Y");
  frame_pull->GetXaxis()->SetTitle("");
  frame_pull->GetYaxis()->SetTitle("");

  TCanvas *c_DsPiMass = new TCanvas("c_DsPiMass","c_DsPiMass",950,900);
  c_DsPiMass->Divide(1,2);

  c_DsPiMass->cd(1);
  frame->Draw();
  TObject *DsPiMassobj = frame->findObject("modelDsPiMass"); 
  TObject *NRobj = frame->findObject("modelNRpdf");
  TObject *D1obj = frame->findObject("modelD1pdf");
  TObject *D2sobj = frame->findObject("modelD2spdf");
  TObject *D1brobj = frame->findObject("modelD1brpdf");
  TLegend *leg = new TLegend(0.85,0.85,1.0,1.0);
  leg->AddEntry(DsPiMassobj,"Total","L");  
  leg->AddEntry(NRobj,"Non-resonant","L");
  leg->AddEntry(D1obj,"D_{1}(2420)^{0}","L");
  leg->AddEntry(D2sobj,"D_{2}^{*}(2460)^{0}","L");
  leg->AddEntry(D1brobj,"D_{1}(2430)^{0}","L");
  leg->Draw("SAME");
  lhcbNameR->Draw("SAME");
  TVirtualPad *pad_mass = c_DsPiMass->GetPad(1);
  pad_mass->SetPad(0.02,0.2,0.98,0.98);

  c_DsPiMass->cd(2);
  TVirtualPad *pad_pull= c_DsPiMass->GetPad(2);
  pad_pull->SetPad(0.02,0.02,0.98,0.2);
  frame_pull->Draw();

  //Summarize results
  r_DsPiMass->Print();
  r_DsPiMass->correlationMatrix().Print();
  cout << "Min - log(L): " << r_DsPiMass->minNll() << "\n";
  cout << "EDM: " << r_DsPiMass->edm() << "\n";

  //Save results
  /*TFile *g = new TFile(outputfile,"RECREATE");
  g->cd();
  c_comp->Write();
  c_DsPiMass->Write();
  r_DsPiMass->Write();*/
}