%%cpp -d
#include <TError.h>
#include <TMath.h>
#include <TCanvas.h>
#include <TRandom3.h>
#include <TFitter.h>
#include <TF1.h>
#include <TStyle.h>
#include <TVector.h>
#include <TGraph.h>

#include "TUnfoldDensity.h"


TRandom *rnd=nullptr;

TH2 *gHistInvEMatrix;

TVirtualFitter *gFitter=nullptr;

%%cpp -d
void chisquare_corr(Int_t &npar, Double_t * /*gin */, Double_t &f, Double_t *u, Int_t /*flag */) {
  //  Minimization function for H1s using a Chisquare method
  //  only one-dimensional histograms are supported
  //  Correlated errors are taken from an external inverse covariance matrix
  //  stored in a 2-dimensional histogram
  Double_t x;
  TH1 *hfit = (TH1*)gFitter->GetObjectFit();
  TF1 *f1   = (TF1*)gFitter->GetUserFunc();

  f1->InitArgs(&x,u);
  npar = f1->GetNpar();
  f = 0;

  Int_t npfit = 0;
  Int_t nPoints=hfit->GetNbinsX();
  Double_t *df=new Double_t[nPoints];
  for (Int_t i=0;i<nPoints;i++) {
    x     = hfit->GetBinCenter(i+1);
    TF1::RejectPoint(kFALSE);
    df[i] = f1->EvalPar(&x,u)-hfit->GetBinContent(i+1);
    if (TF1::RejectedPoint()) df[i]=0.0;
    else npfit++;
  }
  for (Int_t i=0;i<nPoints;i++) {
    for (Int_t j=0;j<nPoints;j++) {
      f += df[i]*df[j]*gHistInvEMatrix->GetBinContent(i+1,j+1);
    }
  }
  delete[] df;
  f1->SetNumberFitPoints(npfit);
}

%%cpp -d
Double_t bw_func(Double_t *x,Double_t *par) {
  Double_t dm=x[0]-par[1];
  return par[0]/(dm*dm+par[2]*par[2]);
}

%%cpp -d
Double_t GenerateEvent(Double_t bgr, // relative fraction of background
                       Double_t mass, // peak position
                       Double_t gamma) // peak width
{
  Double_t t;
  if(rnd->Rndm()>bgr) {
    // generate signal event
    // with positive mass
    do {
      do {
        t=rnd->Rndm();
      } while(t>=1.0);
      t=TMath::Tan((t-0.5)*TMath::Pi())*gamma+mass;
    } while(t<=0.0);
    return t;
  } else {
    // generate background event
    // generate events following a power-law distribution
    //   f(E) = K * TMath::power((E0+E),N0)
    static Double_t const E0=2.4;
    static Double_t const N0=2.9;
    do {
      do {
        t=rnd->Rndm();
      } while(t>=1.0);
      // the mass is returned negative
      // In our example a convenient way to indicate it is a background event.
      t= -(TMath::Power(1.-t,1./(1.-N0))-1.0)*E0;
    } while(t>=0.0);
    return t;
  }
}

%%cpp -d
Double_t DetectorEvent(Double_t mTrue) {
  // smear by double-gaussian
  static Double_t frac=0.1;
  static Double_t wideBias=0.03;
  static Double_t wideSigma=0.5;
  static Double_t smallBias=0.0;
  static Double_t smallSigma=0.1;
  if(rnd->Rndm()>frac) {
    return rnd->Gaus(mTrue+smallBias,smallSigma);
  } else {
    return rnd->Gaus(mTrue+wideBias,wideSigma);
  }
}

TH1::SetDefaultSumw2();

gStyle->SetOptFit(1111);

rnd=new TRandom3();

Double_t const luminosityData=100000;
Double_t const luminosityMC=1000000;
Double_t const crossSection=1.0;

Int_t const nDet=250;
Int_t const nGen=100;
Double_t const xminDet=0.0;
Double_t const xmaxDet=10.0;
Double_t const xminGen=0.0;
Double_t const xmaxGen=10.0;

TH1D *histMgenMC=new TH1D("MgenMC",";mass(gen)",nGen,xminGen,xmaxGen);
TH1D *histMdetMC=new TH1D("MdetMC",";mass(det)",nDet,xminDet,xmaxDet);
TH2D *histMdetGenMC=new TH2D("MdetgenMC",";mass(det);mass(gen)",
                             nDet,xminDet,xmaxDet,nGen,xminGen,xmaxGen);
Int_t neventMC=rnd->Poisson(luminosityMC*crossSection);
for(Int_t i=0;i<neventMC;i++) {
  Double_t mGen=GenerateEvent(0.3, // relative fraction of background
                              4.0, // peak position in MC
                              0.2); // peak width in MC
  Double_t mDet=DetectorEvent(TMath::Abs(mGen));
  // the generated mass is negative for background
  // and positive for signal
  // so it will be filled in the underflow bin
  // this is very convenient for the unfolding:
  // the unfolded result will contain the number of background
  // events in the underflow bin

  // generated MC distribution (for comparison only)
  histMgenMC->Fill(mGen,luminosityData/luminosityMC);
  // reconstructed MC distribution (for comparison only)
  histMdetMC->Fill(mDet,luminosityData/luminosityMC);

  // matrix describing how the generator input migrates to the
  // reconstructed level. Unfolding input.
  // NOTE on underflow/overflow bins:
  //  (1) the detector level under/overflow bins are used for
  //       normalisation ("efficiency" correction)
  //       in our toy example, these bins are populated from tails
  //       of the initial MC distribution.
  //  (2) the generator level underflow/overflow bins are
  //       unfolded. In this example:
  //       underflow bin: background events reconstructed in the detector
  //       overflow bin: signal events generated at masses > xmaxDet
  // for the unfolded result these bins will be filled
  //  -> the background normalisation will be contained in the underflow bin
  histMdetGenMC->Fill(mDet,mGen,luminosityData/luminosityMC);
}

TH2D *histMdetGenSysMC=new TH2D("MdetgenSysMC",";mass(det);mass(gen)",
                                nDet,xminDet,xmaxDet,nGen,xminGen,xmaxGen);
neventMC=rnd->Poisson(luminosityMC*crossSection);
for(Int_t i=0;i<neventMC;i++) {
  Double_t mGen=GenerateEvent
     (0.5, // relative fraction of background
      3.6, // peak position in MC with systematic shift
      0.15); // peak width in MC
  Double_t mDet=DetectorEvent(TMath::Abs(mGen));
  histMdetGenSysMC->Fill(mDet,mGen,luminosityData/luminosityMC);
}

TH1D *histMgenData=new TH1D("MgenData",";mass(gen)",nGen,xminGen,xmaxGen);
TH1D *histMdetData=new TH1D("MdetData",";mass(det)",nDet,xminDet,xmaxDet);
Int_t neventData=rnd->Poisson(luminosityData*crossSection);
for(Int_t i=0;i<neventData;i++) {
  Double_t mGen=GenerateEvent(0.4, // relative fraction of background
                              3.8, // peak position in data
                              0.15); // peak width in data
  Double_t mDet=DetectorEvent(TMath::Abs(mGen));
  // generated data mass for comparison plots
  // for real data, we do not have this histogram
  histMgenData->Fill(mGen);

  // reconstructed mass, unfolding input
  histMdetData->Fill(mDet);
}

TH1D *histDensityGenData=new TH1D("DensityGenData",";mass(gen)",
                                  nGen,xminGen,xmaxGen);
TH1D *histDensityGenMC=new TH1D("DensityGenMC",";mass(gen)",
                                  nGen,xminGen,xmaxGen);
for(Int_t i=1;i<=nGen;i++) {
   histDensityGenData->SetBinContent(i,histMgenData->GetBinContent(i)/
                                     histMgenData->GetBinWidth(i));
   histDensityGenMC->SetBinContent(i,histMgenMC->GetBinContent(i)/
                                     histMgenMC->GetBinWidth(i));
}

TUnfoldDensity unfold(histMdetGenMC,TUnfold::kHistMapOutputVert);

if(unfold.SetInput(histMdetData)>=10000) {
  std::cout<<"Unfolding result may be wrong\n";
}

Int_t nScan=30;

Double_t tauMin=0.0;
Double_t tauMax=0.0;
Int_t iBest;
TSpline *logTauX,*logTauY;
TGraph *lCurve;

#ifdef VERBOSE_LCURVE_SCAN
Int_t oldinfo=gErrorIgnoreLevel;
gErrorIgnoreLevel=kInfo;
#endif

iBest=unfold.ScanLcurve(nScan,tauMin,tauMax,&lCurve,&logTauX,&logTauY);

#ifdef VERBOSE_LCURVE_SCAN
gErrorIgnoreLevel=oldinfo;
#endif

Double_t SYS_ERROR1_MSTART=6;
Double_t SYS_ERROR1_SIZE=0.1;
TH2D *histMdetGenSys1=new TH2D("Mdetgensys1",";mass(det);mass(gen)",
                               nDet,xminDet,xmaxDet,nGen,xminGen,xmaxGen);
for(Int_t i=0;i<=nDet+1;i++) {
   if(histMdetData->GetBinCenter(i)>=SYS_ERROR1_MSTART) {
      for(Int_t j=0;j<=nGen+1;j++) {
         histMdetGenSys1->SetBinContent(i,j,SYS_ERROR1_SIZE);
      }
   }
}
unfold.AddSysError(histMdetGenSysMC,"SYSERROR_MC",TUnfold::kHistMapOutputVert,
                   TUnfoldSys::kSysErrModeMatrix);
unfold.AddSysError(histMdetGenSys1,"SYSERROR1",TUnfold::kHistMapOutputVert,
                   TUnfoldSys::kSysErrModeRelative);

std::cout<<"tau="<<unfold.GetTau()<<"\n";
std::cout<<"chi**2="<<unfold.GetChi2A()<<"+"<<unfold.GetChi2L()
         <<" / "<<unfold.GetNdf()<<"\n";
std::cout<<"chi**2(sys)="<<unfold.GetChi2Sys()<<"\n";

Double_t t[1],x[1],y[1];
logTauX->GetKnot(iBest,t[0],x[0]);
logTauY->GetKnot(iBest,t[0],y[0]);
TGraph *bestLcurve=new TGraph(1,x,y);
TGraph *bestLogTauLogChi2=new TGraph(1,t,x);

TH1 *histMunfold=unfold.GetOutput("Unfolded");

TH1 *histMdetFold=unfold.GetFoldedOutput("FoldedBack");

TH2 *histEmatTotal=unfold.GetEmatrixTotal("EmatTotal");

TH1D *histTotalError=
   new TH1D("TotalError",";mass(gen)",nGen,xminGen,xmaxGen);
for(Int_t bin=1;bin<=nGen;bin++) {
  histTotalError->SetBinContent(bin,histMunfold->GetBinContent(bin));
  histTotalError->SetBinError
     (bin,TMath::Sqrt(histEmatTotal->GetBinContent(bin,bin)));
}

TH1 *histRhoi=unfold.GetRhoItotal("rho_I",
                                  nullptr, // use default title
                                  nullptr, // all distributions
                                  "*[UO]", // discard underflow and overflow bins on all axes
                                  kTRUE, // use original binning
                                  &gHistInvEMatrix // store inverse of error matrix
                                  );

gFitter=TVirtualFitter::Fitter(histMunfold);
gFitter->SetFCN(chisquare_corr);

TF1 *bw=new TF1("bw",bw_func,xminGen,xmaxGen,3);
bw->SetParameter(0,1000.);
bw->SetParameter(1,3.8);
bw->SetParameter(2,0.2);

histMunfold->Fit(bw,"UE");

TCanvas output;
output.Divide(3,2);

output.cd(1);
histMdetGenMC->Draw("BOX");

output.cd(2);
histTotalError->SetLineColor(kBlue);
histTotalError->Draw("E");
histMunfold->SetLineColor(kGreen);
histMunfold->Draw("SAME E1");
histDensityGenData->SetLineColor(kRed);
histDensityGenData->Draw("SAME");
histDensityGenMC->Draw("SAME HIST");

output.cd(3);
histMdetFold->SetLineColor(kBlue);
histMdetFold->Draw();
histMdetMC->Draw("SAME HIST");

TH1 *histInput=unfold.GetInput("Minput",";mass(det)");

histInput->SetLineColor(kRed);
histInput->Draw("SAME");

output.cd(4);
histRhoi->Draw();

output.cd(5);
logTauX->Draw();
bestLogTauLogChi2->SetMarkerColor(kRed);
bestLogTauLogChi2->Draw("*");

output.cd(6);
lCurve->Draw("AL");
bestLcurve->SetMarkerColor(kRed);
bestLcurve->Draw("*");

output.SaveAs("testUnfold1.ps");

return 0;

gROOT->GetListOfCanvases()->Draw()