#include "RooConstVar.h"
#include "RooStepFunction.h"
#include "TF1.h"
#include "TFitResult.h"
#include "TGraphErrors.h"
#include <RooDataHist.h>
#include <RooFit.h>
#include <RooHelpers.h>
#include <RooHistPdf.h>
#include <RooMomentMorph.h>
#include <RooMomentMorphND.h>
#include <RooGenericPdf.h>
#include <RooParametricStepFunction.h>
#include <RooWorkspace.h>
#include <RooWrapperPdf.h>
#include <TFile.h>
#include <TH3D.h>
#include <TStyle.h>
#include <TSystemDirectory.h>
#include <iostream>

using namespace RooFit;
using namespace RooStats;
using namespace std;


RooHistPdf* loadHistSnapshot(RooAbsPdf* pdf,const RooRealVar* x,const RooRealVar* y,const RooArgSet* obs,TString unique="",int interpOrder=0,double* norm=nullptr){

  auto hist = (TH2D*) pdf->createHistogram(unique,*x,YVar(*y));
  TString name = pdf->GetName();
  name+=unique;

  if (norm!=nullptr){
    *norm=hist->Integral();
  }
  
  auto datahist = new RooDataHist(Form("dataHist%s",name.Data()),
				  Form("dataHist%s",name.Data()),*obs,
				  hist);
  auto histpdf = new RooHistPdf(Form("pdf_%s",name.Data()),
				Form("pdf_%s",name.Data()),
				*obs,*datahist,interpOrder);
  
  return histpdf;
}


void loadGaussians(RooWorkspace* w,TString pdfname,int npdfs=10,int axis=0){

  RooRealVar* morphVar;
  if (axis==0){
    morphVar = w->var("par_X");
  }
  else{
    morphVar = w->var("par_Y");
  }

  double min = morphVar->getMin();
  double max = morphVar->getMax();
  
  auto pdfList = new RooArgList("pdfList");
  auto morphValList = new RooArgList("morphValList");
  
  
  double currentVal;
  vector<pair<double,double>> normMap; // form is morphParamValue:normalizationOfPdfAtMorphParamValue
  double stepSize = TMath::Log(max/(min+1));
  for (int i=0; i<=npdfs;i++){ //= include the max val
    currentVal = (min+1)*TMath::Exp(stepSize*i/npdfs);
    morphVar->setVal(currentVal);
    double norm;
    auto histi = loadHistSnapshot(w->pdf(pdfname),w->var("X"),w->var("Y"),w->set("obs"),Form("_%g",currentVal),0,&norm);
    normMap.push_back(make_pair(currentVal,norm));
    //w->import(*histi);
    pdfList->add(*histi);
    morphValList->add(RooConst(currentVal));
  }


  auto coefList = new RooArgList("norms_gaus");
  auto limits = new RooArgList("limits_gaus");
  
  sort(normMap.begin(),normMap.end()); // sort by first (morph param value)
  // we need to add the first point twice to extend to the minimum ( I think / it makes the plots look nicer at the edges)
  coefList->add(RooConst(normMap.front().second));
  limits->add(RooConst(-RooNumber::infinity()));
  
  limits->add(RooConst(morphVar->getMin()));
  for (auto pair: normMap){
    cout << "normMap = "<< pair.first << "\t" << pair.second << endl;
      
    coefList->add(RooConst(pair.second));
    limits->add(RooConst(pair.first));
  }
  // we need to add that last point twice so it can extend to the maximum (think about this in the binned case)
  coefList->add(RooConst(normMap.back().second));
  limits->add(RooConst(RooNumber::infinity()));
  auto rate = new RooStepFunction("rate_gaus","rate_gaus",*morphVar,*coefList,*limits,kTRUE); // last bool turns on or off interpolation!
  // rate is a function not a PDF here!

  w->import(*rate);
  
  // set number of bin in shape parameter to be
  // larger than the number of templates. this way
  // we avoid having to run a ton of MC
  morphVar->setBins(2*npdfs);

  // make the MomentMorph (MM) objects - with and without horizontal morphing for testing
  vector<RooMomentMorph::Setting> settings = { RooMomentMorph::Linear};//, RooMomentMorph::NonLinearPosFractions, RooMomentMorph::NonLinearLinFractions, RooMomentMorph::SineLinear };
  vector<TString> names = {"lin"};//,"nonlinpos","nonlinlin","sinelin"};
  for (int i=0;i< names.size();i++){
    auto mm = new RooMomentMorph(Form("mm2_%d_%s",axis,names[i].Data()),Form("mm_%d_%s",axis,names[i].Data()),*morphVar,*w->set("obs"),*pdfList,*morphValList,settings[i]);
    
    w->import(*mm,RecycleConflictNodes()); // RecycleConflictNodes tells the importer to reuse all the definitions of objects that already exist in the workspace
    // this essentially tells it that all the MMs in this loop use the same set of pdfs (otherwise it freaks out that they all have the same name)

    auto dh = new RooDataHist(Form("dh_mm2_%d_%s",axis,names[i].Data()),Form("dh_mm2_%d_%s",axis,names[i].Data()),RooArgSet(*w->set("obs"),morphVar)); // alright. this last variable added to the RooArgSet needs to be a pointer to a RooRealVar and not the acutal object. Passing an object takes O(10x) longer to generate hists!
    // must be somethign about the set owning its contents or not
    mm->fillDataHist(dh,w->set("obs"),1,kTRUE,kTRUE);

    auto dhpdf = new RooHistPdf(Form("pdf_mm2_%d_%s",axis,names[i].Data()),Form("pdf_mm2_%d_%s",axis,names[i].Data()),RooArgSet(*w->set("obs"),morphVar),*dh);
    auto weight = new RooGenericPdf("test_weight","test weight","2*x[0]",*morphVar);
    auto wdh = new RooProdPdf(Form("weighted_dh_%d_%s",axis,names[i].Data()),Form("weighted_dh_%d_%s",axis,names[i].Data()),*weight,Conditional(*dhpdf,*w->set("obs")));
    w->import(*wdh,RecycleConflictNodes());
    auto hist = (TH3D*) wdh->createHistogram(Form("hist_mm2_%d_%s",axis,names[i].Data()),*w->var("X"),RooFit::YVar(*w->var("Y")),RooFit::ZVar(*morphVar));
    auto c = new TCanvas();
    hist->Draw("ISO");
    c->SaveAs(Form("mm2_%d_%s_3Dinterp.png",axis,names[i].Data()));
 
    
  }
  
  
  return;

}

void plot_nll(RooWorkspace* w, TString morphvarname="par_X", TString pdfname="pdf_mm2_0_lin", int n_evts=100){
   
   w->var(morphvarname.Data())->setVal(30);
   auto ds = w->pdf(pdfname.Data())->generate(*w->set("obs"),n_evts);
      w->var(morphvarname.Data())->setVal(10);
   auto frame = w->var(morphvarname.Data())->frame();
   
   auto nll = w->pdf(pdfname.Data())->createNLL(*ds);
   
   auto c = new TCanvas();
   nll->plotOn(frame);
   frame->Draw();
   c->SaveAs(Form("nll_%s_%devts.png",pdfname.Data(),n_evts));
   
   
   // do a fit as well
   w->pdf(pdfname.Data())->fitTo(*ds,RooFit::Minimizer("Minuit","Migrad"),RooFit::PrintLevel(1));
   //auto c2 = new TCanvas();
   //plotDataset(w->pdf(pdfname.Data()),w->set("obs"),nullptr,ds,"X,Y","","",c2,"colz");
   //c2->SaveAs(Form("fit_res_%devts.png",n_evts));
   
   
   return;
}

void gauss_example(){

  auto w = new RooWorkspace("w");
  w->factory("X[10,0,80]"); //initial_value,min,max
  w->var("X")->setBins(80);
  w->factory("Y[4,2.5,5.5]");
  w->var("Y")->setBins(80);
  w->defineSet("obs","X,Y");



  int nSlices = 10;
  w->factory("PROD::gaus2d(Gaussian::gausx(X,par_X[10,0,80],5),Gaussian::gausy(Y,par_Y[4,2.5,5.5],0.5))");
  loadGaussians(w,"gaus2d",nSlices,0);
  //w->var("X")->setVal(40.);
  //loadGaussians(w,"gaus2d",nSlices,1);

   double maxV = 80;
   double minV=0;
   double currentVal;
   double stepSize = TMath::Log(maxV/(minV+1));
   TCanvas *c[nSlices];
   for (int i=0; i<=nSlices;i++){
     currentVal = (TMath::Exp(stepSize*i/nSlices))*(minV+1);
     Int_t lineColor = gStyle->GetColorPalette(int(255. * (i / double(nSlices + 1))));
     gStyle->SetHistLineColor(lineColor);
     //plotDataset(w->pdf(Form("pdf_gaus2d_%g",currentVal)),w->set("obs"),nullptr,nullptr,"X,Y","","",nullptr,"colz");
     c[i] = new TCanvas();
     auto h = w->pdf(Form("pdf_gaus2d_%g",currentVal))->createHistogram("X,Y");
     h->Draw("colz");
     gPad->SaveAs(Form("gaus_template_%g.png",currentVal));
   }
   //gPad->SaveAs("gaus_templates.png");
  
  plot_nll(w, TString("par_X"), TString("pdf_mm2_0_lin"), 1000);
  
  w->SaveAs("gaus_morph_test_workspace.root");
  

  return;
}

int main(){
    gauss_example(); // allow for both compiled and quick command line calls to this script
    return 1;
}