#include "Riostream.h"

// Custom exponential PDF
class ExpPdf : public RooAbsPdf {
public:
  ExpPdf() {};
  ExpPdf(const char *name, const char *title, RooAbsReal& _t, RooAbsReal& _tau);
  ExpPdf(const ExpPdf& other, const char* name = 0);
  virtual TObject* clone(const char* newname) const { return new ExpPdf(*this, newname); }
  inline virtual ~ExpPdf() { }

  Int_t getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* rangeName=0) const ;
  Double_t analyticalIntegral(Int_t code, const char* rangeName=0) const ;
  Double_t indefiniteIntegral(Double_t y) const;

protected:
  RooRealProxy t;
  RooRealProxy tau;
  Double_t evaluate() const;
};

ExpPdf::ExpPdf(const char *name, const char *title, RooAbsReal& _t, RooAbsReal& _tau) :
  RooAbsPdf(name, title), t("t", "t", this, _t), tau("tau", "tau", this, _tau) {
}

ExpPdf::ExpPdf(const ExpPdf& other, const char* name) : RooAbsPdf(other, name), t("t", this, other.t), tau("tau", this, other.tau){
}

Double_t ExpPdf::evaluate() const {
  if (t < 0) return 0.;
  return exp(-t/tau);
}

Double_t ExpPdf::indefiniteIntegral(Double_t y) const {
  return -exp(-y/tau)*tau;
}

Int_t ExpPdf::getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* rangeName) const {
  if (matchArgs(allVars,analVars,t)) return 1;
  return 0;
}

Double_t ExpPdf::analyticalIntegral(Int_t code, const char* rangeName) const {
  if (code == 1){
    Double_t t1 = t.min(rangeName);
    Double_t t2 = t.max(rangeName);
    if (t2 <= 0) return 0;
    t1 = TMath::Max(0.,t1);
    return indefiniteIntegral(t2) - indefiniteIntegral(t1);
  }
  return 0;
}

// Main macro
void conv_linearity_test() {

  // Create observable
  RooRealVar* t = new RooRealVar("t", "time axis", 0, 10, "ns");
  t->setBins(100);

  // Coefficient to convert fwhm to dispersion
  RooConstVar fwhm2dispersion = RooFit::RooConst(0.424661);

  // FIRST TEST: RooAddPdf(ExpPdf + ExpPdf) x gauss
  // Resolution function is RooGaussian
  RooRealVar* mean1 = new RooRealVar("mean1", "gauss1 mean", 3, 0, 10, "ns");
  RooRealVar* fwhm1 = new RooRealVar("fwhm1", "gauss1 fwhm", 0.5, 0, 2, "ns");
  RooFormulaVar* sigma1 = new RooFormulaVar("sigma1", "gauss1 dispersion", "@0*@1", RooArgList(*fwhm1, fwhm2dispersion));
  RooGaussian* gauss1 = new RooGaussian("gauss1", "gauss pdf", *t, *mean1, *sigma1);

  // Model is sum of two custom ExpPdf's. Their sum is convoluted via RooFFTConvPdf.
  RooRealVar* tauA1 = new RooRealVar("tauA1", "lifetime A1", 0.2, "ns");
  RooRealVar* tauB1 = new RooRealVar("tauB1", "lifetime B1", 1, "ns");
  RooRealVar* I_B1 = new RooRealVar("I_B1", "intensity of B1", 0.13, 0.0, 1.0);
  ExpPdf* pdfA1 = new ExpPdf("pdfA1", "PDF A1", *t, *tauA1);
  ExpPdf* pdfB1 = new ExpPdf("pdfB1", "PDF B1", *t, *tauB1);
  RooAddPdf* pdfA1B1 = new RooAddPdf("pdfAB", "sum pdfA1 and pdf B1", RooArgList(*pdfB1, *pdfA1), *I_B1);   // first sum
  // Uncomment this line to get the correct result.
  // pdfA1B1->fixAddCoefNormalization(RooArgSet(*t));
  RooFFTConvPdf* pdf1 = new RooFFTConvPdf("pdf1", "RooAddPdf(ExpPdf + ExpPdf) x gauss", *t, *pdfA1B1, *gauss1); // then convolute

  // SECOND TEST: RooAddPdf((ExpPdf x gauss) + (ExpPdf x gauss))
  // Resolution function is RooGaussian with same parameter starting values
  RooRealVar* mean2 = new RooRealVar("mean2", "gauss2 mean", 3, 0, 10, "ns");
  RooRealVar* fwhm2 = new RooRealVar("fwhm2", "gauss2 fwhm", 0.5, 0, 2, "ns");
  RooFormulaVar* sigma2 = new RooFormulaVar("sigma2", "gauss2 dispersion", "@0*@1", RooArgList(*fwhm2, fwhm2dispersion));
  RooGaussian* gauss2 = new RooGaussian("gauss2", "gauss pdf", *t, *mean2, *sigma2);

  // Model is sum of two custom ExpPdf's independently convoluted via RooFFTConvPdf
  RooRealVar* tauA2 = new RooRealVar("tauA2", "lifetime A2", 0.2, "ns");
  RooRealVar* tauB2 = new RooRealVar("tauB2", "lifetime B2", 1, "ns");
  RooRealVar* I_B2 = new RooRealVar("I_B2", "intensity of B2", 0.13, 0.0, 1.0);
  ExpPdf* pdfA2 = new ExpPdf("pdfA2", "PDF A2", *t, *tauA2);
  RooFFTConvPdf* convA2 = new RooFFTConvPdf("convA2", "convoluted A2", *t, *pdfA2, *gauss2);
  ExpPdf* pdfB2 = new ExpPdf("pdfB2", "PDF B2", *t, *tauB2);
  RooFFTConvPdf* convB2 = new RooFFTConvPdf("convB2", "convoluted B2", *t, *pdfB2, *gauss2);
  RooAddPdf* pdf2 = new RooAddPdf("pdf2", "RooAddPdf((ExpPdf x gauss) + (ExpPdf x gauss))", RooArgList(*convB2, *convA2), *I_B2);

  // Create dataset from second model (that gives correct intensities)
  RooDataHist* data = pdf2->generateBinned(*t, 10000);

  // FIT DATA (to the identical dataset)
  pdf1->fitTo(*data);
  pdf2->fitTo(*data);

  // PLOT FITS AND REVEAL THE PROBLEM:
  TCanvas* canvas = new TCanvas("canvas", "canvas", 1280, 800);
  canvas->Divide(1, 2);

  // Top plot: convolution of two RooDecay's and RooGaussModel
  canvas->cd(1)->SetLogy();
  RooPlot* frame1 = t->frame(RooFit::Title("RooAddPdf(ExpPdf + ExpPdf) x gauss, gives incorrect I_B1 = 0.03 unless I call fixAddCoefNormalization()"));
  data->plotOn(frame1, RooFit::MarkerSize(0.2));
  pdf1->plotOn(frame1);
  gauss1->plotOn(frame1, RooFit::LineStyle(kDashed));
  pdf1->paramOn(frame1, RooFit::Layout(0.58, 0.9, 0.9), RooFit::Format("NEU", RooFit::AutoPrecision(1)));
  frame1->Draw();

  // Middle plot: convolution of custom ExpPdf's with RooGaussian by means of RooFFTConvPdf
  canvas->cd(2)->SetLogy();
  RooPlot* frame2 = t->frame(RooFit::Title("RooAddPdf((ExpPdf x gauss) + (ExpPdf x gauss)), gives correct I_B2 = 0.136"));
  data->plotOn(frame2, RooFit::MarkerSize(0.2));
  pdf2->plotOn(frame2);
  gauss2->plotOn(frame2, RooFit::LineStyle(kDashed));
  pdf2->paramOn(frame2, RooFit::Layout(0.58, 0.9, 0.9), RooFit::Format("NEU", RooFit::AutoPrecision(1)));
  frame2->Draw();
}