%%cpp -d

#include <iostream>

#include "TROOT.h"
#include "TSystem.h"
#include "TStyle.h"
#include "TRandom3.h"
#include "TString.h"
#include "TMath.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TLegend.h"
#include "TCanvas.h"
#include "TColor.h"
#include "TLine.h"

#include "TSVDUnfold.h"


Double_t Reconstruct( Double_t xt, TRandom3& R )
{
   // apply some Gaussian smearing + bias and efficiency corrections to fake reconstruction
   const Double_t cutdummy = -99999.0;
   Double_t xeff = 0.3 + (1.0 - 0.3)/20.0*(xt + 10.0);  // efficiency
   Double_t x    = R.Rndm();
   if (x > xeff) return cutdummy;
   else {
     Double_t xsmear= R.Gaus(-2.5,0.2); // bias and smear
     return xt+xsmear;
   }
}

gROOT->SetStyle("Plain");
gStyle->SetOptStat(0);

TRandom3 R;

const Double_t cutdummy= -99999.0;

Int_t nbins = 40;
TH1D *xini = new TH1D("xini", "MC truth", nbins, -10.0, 10.0);
TH1D *bini = new TH1D("bini", "MC reco", nbins, -10.0, 10.0);
TH2D *Adet = new TH2D("Adet", "detector response", nbins, -10.0, 10.0, nbins, -10.0, 10.0);

TH1D *data = new TH1D("data", "data", nbins, -10.0, 10.0);

TH1D *datatrue = new TH1D("datatrue", "data truth", nbins, -10.0, 10.0);

TH2D *statcov = new TH2D("statcov", "covariance matrix", nbins, -10.0, 10.0, nbins, -10.0, 10.0);

for (Int_t i= 0; i<100000; i++) {
   Double_t xt = R.BreitWigner(0.3, 2.5);
   xini->Fill(xt);
   Double_t x = Reconstruct( xt, R );
   if (x != cutdummy) {
      Adet->Fill(x, xt);
      bini->Fill(x);
   }
}

for (Int_t i=0; i<10000; i++) {
   Double_t xt = R.Gaus(0.0, 2.0);
   datatrue->Fill(xt);
   Double_t x = Reconstruct( xt, R );
   if (x != cutdummy)
   data->Fill(x);
}

cout << "Created toy distributions and errors for: " << endl;
cout << "... \"true MC\"   and \"reconstructed (smeared) MC\"" << endl;
cout << "... \"true data\" and \"reconstructed (smeared) data\"" << endl;
cout << "... the \"detector response matrix\"" << endl;

for (int i=1; i<=data->GetNbinsX(); i++) {
    statcov->SetBinContent(i,i,data->GetBinError(i)*data->GetBinError(i));
}

TSVDUnfold *tsvdunf = new TSVDUnfold( data, statcov, bini, xini, Adet );

tsvdunf->SetNormalize( kFALSE ); // no normalisation here

TH1D* unfres = tsvdunf->Unfold( 13 );

TH1D* ddist = tsvdunf->GetD();

TH1D* svdist = tsvdunf->GetSV();

TH2D* ustatcov = tsvdunf->GetUnfoldCovMatrix( statcov, 100 );

TH2D* uadetcov = tsvdunf->GetAdetCovMatrix( 100 );

ustatcov->Add( uadetcov );

TH2D* utaucov = tsvdunf->GetXtau();
utaucov->Add( uadetcov );

TH2D* uinvcov = tsvdunf->GetXinv();

for (int i=1; i<=unfres->GetNbinsX(); i++) {
   unfres->SetBinError(i, TMath::Sqrt(utaucov->GetBinContent(i,i)));
}

xini->Scale(0.7*datatrue->Integral()/xini->Integral());

TLegend *leg = new TLegend(0.58,0.60,0.99,0.88);
leg->SetBorderSize(0);
leg->SetFillColor(0);
leg->SetFillStyle(0);
leg->AddEntry(unfres,"Unfolded Data","p");
leg->AddEntry(datatrue,"True Data","l");
leg->AddEntry(data,"Reconstructed Data","l");
leg->AddEntry(xini,"True MC","l");

TCanvas *c1 = new TCanvas( "c1", "Unfolding toy example with TSVDUnfold", 1000, 900 );

c1->Divide(1,2);
TVirtualPad * c11 = c1->cd(1);

TH1D* frame = new TH1D( *unfres );
frame->SetTitle( "Unfolding toy example with TSVDUnfold" );
frame->GetXaxis()->SetTitle( "x variable" );
frame->GetYaxis()->SetTitle( "Events" );
frame->GetXaxis()->SetTitleOffset( 1.25 );
frame->GetYaxis()->SetTitleOffset( 1.29 );
frame->Draw();

data->SetLineStyle(2);
data->SetLineColor(4);
data->SetLineWidth(2);
unfres->SetMarkerStyle(20);
datatrue->SetLineColor(2);
datatrue->SetLineWidth(2);
xini->SetLineStyle(2);
xini->SetLineColor(8);
xini->SetLineWidth(2);

unfres->Draw("same");
datatrue->Draw("same");
data->Draw("same");
xini->Draw("same");

leg->Draw();

TVirtualPad * c12 = c1->cd(2);
c12->Divide(2,1);
TVirtualPad * c2 = c12->cd(1);
c2->SetRightMargin   ( 0.15         );

TH2D* covframe = new TH2D( *ustatcov );
covframe->SetTitle( "TSVDUnfold covariance matrix" );
covframe->GetXaxis()->SetTitle( "x variable" );
covframe->GetYaxis()->SetTitle( "x variable" );
covframe->GetXaxis()->SetTitleOffset( 1.25 );
covframe->GetYaxis()->SetTitleOffset( 1.29 );
covframe->Draw();

ustatcov->SetLineWidth( 2 );
ustatcov->Draw( "colzsame" );

TVirtualPad * c3 = c12->cd(2);
c3->SetLogy();

TLine *line = new TLine( 0.,1.,40.,1. );
line->SetLineStyle(2);

TH1D* dframe = new TH1D( *ddist );
dframe->SetTitle( "TSVDUnfold |d_{i}|" );
dframe->GetXaxis()->SetTitle( "i" );
dframe->GetYaxis()->SetTitle( "|d_{i}|" );
dframe->GetXaxis()->SetTitleOffset( 1.25 );
dframe->GetYaxis()->SetTitleOffset( 1.29 );
dframe->SetMinimum( 0.001 );
dframe->Draw();

ddist->SetLineWidth( 2 );
ddist->Draw( "same" );
line->Draw();

gROOT->GetListOfCanvases()->Draw()