using namespace RooFit;

#include "RooRealVar.h"
#include "RooDataSet.h"
#include "TGraph.h"
#include "TStyle.h"
#include "TMath.h"
#include "RooAbsPdf.h"
#include "RooGaussian.h"
#include "RooFFTConvPdf.h"
#include "RooNumConvPdf.h"
#include "RooFormulaVar.h"
#include "RooPlot.h"



void test_convolution(int type=0,int readwrite=0)
{
  // define x1 shape
  
  RooRealVar x1("x1", "mass of x1", 400, 1000);
  RooRealVar x1_mean("x1_mean", "mean mass", 780, 760, 850);
  RooRealVar x1_width("x1_width", "width of x1 resonance", 150);

  RooBreitWigner x1_bw("x1_bw", "nonrel bw", x1, x1_mean, x1_width);

  // define gaussian describing resolution of x1

  RooRealVar mg("mg","mg",0) ; 
  RooRealVar sg("sg","sg",35, 10, 60) ; 
  RooGaussian resolution("resolution","gauss",x1,mg,sg) ; 
  


  //-----------------------------------------------------------
  //-----------------------------------------------------------
  // Construct convolution
  RooAbsPdf * pdf_conv = nullptr;
  
  if (type == 0) 
  // 1) use RooVoigtian
   pdf_conv = new RooVoigtian("x1_bw_conv", "Voigtian", x1,  x1_mean, x1_width, sg);

  // 2) use RooFFTConvPdf
  else {
    x1.setBins(10000,"cache") ; 
    pdf_conv = new RooFFTConvPdf ("x1_bw_conv", " BW (X) gauss", x1, x1_bw, resolution);
  }
  //-----------------------------------------------------------
  //-----------------------------------------------------------

  RooAbsPdf & x1_bw_conv = *pdf_conv;


  // define parameters x2
  RooRealVar x2("x2", "mass of x2", 5000, 5600);
  RooRealVar x2_mean("x2_mean", "mean of Bmass from mc",5279);
  RooRealVar x2_sigma("x2_sigma", "sigma of gaussian 1",13, 5, 20);

  RooGaussian gauss("gauss","gauss",x2,x2_mean,x2_sigma);

  // create 2 dimensional pdf for x1 and x2
  RooProdPdf *sig = new RooProdPdf("sig", "signal from Bplus and x1plus", RooArgList(x1_bw_conv, gauss));
  

  // generate dataset from pdfs
 
  RooAbsData *data = nullptr;
  if (readwrite == 1) {
    auto file  = TFile::Open("test_convolution.root");
    auto w = (RooWorkspace*) file->Get("w");
    data = w->data("sigData");
  }
  else 
    data = sig->generate(RooArgSet(x2, x1), 8000);


  // create nll var
  auto nll = sig->createNLL(*data,RooFit::Optimize(0));

  auto pl_sg = sg.frame();
  nll->plotOn(pl_sg);
  new TCanvas();
  //pl_sg->Draw();
  
  // fix other parameters 
  x1_mean.setConstant(true);
  x2_sigma.setConstant(true);
  //RooFitResult * result = sig->fitTo(*data, RooFit::Minimizer("Minuit2"));
  //nll->constOptimizeTestStatistic(RooAbsArg::DeActivate);

  // scan values of sg


#if 0
  auto h1 = new TH1D("h1","h1",50,10,60);
  for (int i = 1; i <= h1->GetNbinsX(); ++i) {
    sg.setVal(h1->GetBinCenter(i) );
    h1->SetBinContent(i,nll->getVal() ); 
  }
  h1->Draw("HIST");


  std::cout << "create histogram" << std::endl;
  auto h2 = nll->createHistogram("sg");
  new TCanvas();
  h2->Draw("HIST");
#endif

  std::cout << "create function using nll->asTF " << std::endl;

  auto f1 = nll->asTF(sg);
  f1->SetLineColor(kRed);
  new TCanvas();
  f1->DrawClone();

  
  std::cout << "create function using create projection " << std::endl;

  auto proj = nll->createPlotProjection(sg, RooArgSet() ); 

  f1 = proj->asTF(sg);
  f1->SetLineColor(kRed);
  new TCanvas();
  f1->DrawClone();

  

  std::cout << "Create minimizer... " << std::endl;

  RooMinimizer minim(*nll);
  minim.setPrintLevel(3);
  minim.minimize("Minuit2");
  
  /////////////////////////////////////////////////////
  // plotting
  // plot result for x1
  TCanvas* canvas1 = new TCanvas(TString::Format("canvas1_%d",type),"canvas1");

  RooPlot * frame1 = x1.frame(Title("X1"));//, Range("full")
  data->plotOn(frame1, Binning(60));
  sig->plotOn(frame1, LineColor(kBlack));
  sig->paramOn(frame1, data);
   
  frame1->Draw();

  // plot result for x2
  TCanvas* canvas2 = new TCanvas(TString::Format("canvas2_%d",type),"canvas2");


  RooPlot * frame22 =  x2.frame(Title("X2"));//, Range("full")
  data->plotOn(frame22, Binning(60));
  sig->plotOn(frame22, LineColor(kBlack));
  sig->paramOn(frame22, data);
   
  frame22->Draw();

  if (readwrite == 2) { 
    RooWorkspace w("w");
    w.import(*data);
    w.writeToFile("test_convolution.root");
  }

}