////////////////////////////////////////////////////////////////////////////
//                                                                        //
// This program should be run under root:                                 //
//      root fluxMUnfold.C++                                              //
//                                                                        //
// Authors: T. Bretz  02/2002 <mailto:tbretz@astro.uni-wuerzburg.de>      //
//          W. Wittek 09/2002 <mailto:wittek@mppmu.mpg.de>                //
//          R. Wagner 11/2004 <mailto:rwagner@mppmu.mpg.de>               //
//                                                                        //
// this macro is prepared to be used in the analysis:                     //
//                                                                        //
//      the unfolding should be called by                                 //
//      doUnfolding(TH2D &tobeunfolded,      // (E-est, Theta)            //
//                  TH3D &migrationmatrix,   // (E-est, E-true, Theta)    //
//                  TH2D &unfolded)          // (E-true,Theta)            //
//                                                                        //
////////////////////////////////////////////////////////////////////////////

#include <TMath.h>
#include <TRandom3.h>
#include <TVector.h>
#include <TMatrixD.h>
#include <TMatrix.h>
#include <TH1.h>
#include <TH2.h>
#include <TH3.h>
#include <TProfile.h>
#include <TF1.h>
#include <iostream.h>
#include <TMinuit.h>
#include <TCanvas.h>
#include <TMarker.h>

#include <fstream.h>
#include <iomanip.h>

TH1 *DrawMatrixClone(const TMatrixD &m, Option_t *opt="")
{
    const Int_t nrows = m.GetNrows();
    const Int_t ncols = m.GetNcols();

    TMatrix m2(nrows, ncols);
    for (int i=0; i<nrows; i++)
        for (int j=0; j<ncols; j++)
            m2(i, j) = m(i, j);

    TH2F *hist = new TH2F(m2);
    hist->SetBit(kCanDelete);
    hist->Draw(opt);
    hist->SetDirectory(NULL);

    return hist;

}

TH1 *DrawMatrixColClone(const TMatrixD &m, Option_t *opt="", Int_t col=0)
{
    const Int_t nrows = m.GetNrows();

    TVector vec(nrows);
    for (int i=0; i<nrows; i++)
        vec(i) = m(i, col);

    TH1F *hist = new TH1F("TVector","",nrows,0,nrows);
    for (int i=0; i<nrows; i++)
    {
      hist->SetBinContent(i+1, vec(i));
    }

    hist->SetBit(kCanDelete);
    hist->Draw(opt);
    hist->SetDirectory(NULL);

    return hist;
}


void PrintTH3Content(const TH3 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
    {
      for (Int_t j=1; j<=hist.GetNbinsY(); j++)
        for (Int_t k=1; k<=hist.GetNbinsZ(); k++)
            cout << hist.GetBinContent(i,j,k) << " \t";
      cout << endl << endl;
    }
}

void PrintTH3Error(const TH3 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << " <error>" << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
    {
      for (Int_t j=1; j<=hist.GetNbinsY(); j++)
        for (Int_t k=1; k<=hist.GetNbinsZ(); k++)
            cout << hist.GetBinError(i, j, k) << " \t";
      cout << endl << endl;
    }
}

void PrintTH2Content(const TH2 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
        for (Int_t j=1; j<=hist.GetNbinsY(); j++)
            cout << hist.GetBinContent(i,j) << " \t";
        cout << endl << endl;
}

void PrintTH2Error(const TH2 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << " <error>" << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
        for (Int_t j=1; j<=hist.GetNbinsY(); j++)
            cout << hist.GetBinError(i, j) << " \t";
        cout << endl << endl;
}

void PrintTH1Content(const TH1 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
        cout << hist.GetBinContent(i) << " \t";
    cout << endl << endl;
}

void PrintTH1Error(const TH1 &hist)
{
    cout << hist.GetName() << ": " << hist.GetTitle() << " <error>" << endl;
    cout << "-----------------------------------------------------" << endl;
    for (Int_t i=1; i<=hist.GetNbinsX(); i++)
        cout << hist.GetBinError(i) << " \t";
    cout << endl << endl;
}

void CopyCol(TMatrixD &m, const TH1 &h, Int_t col=0)
{
    const Int_t n = m.GetNrows();

    for (Int_t i=0; i<n; i++)
        m(i, col) = h.GetBinContent(i+1);
}

void CopyCol(TH1 &h, const TMatrixD &m, Int_t col=0)
{
    const Int_t n = m.GetNrows();

    for (Int_t i=0; i<n; i++)
        h.SetBinContent(i+1, m(i, col));
}

void CopyH2M(TMatrixD &m, const TH2 &h)
{
    const Int_t nx = m.GetNrows();
    const Int_t ny = m.GetNcols();

    for (Int_t i=0; i<nx; i++)
        for (Int_t j=0; j<ny; j++)
            m(i, j) = h.GetBinContent(i+1, j+1);
}

void CopySqr(TMatrixD &m, const TH1 &h)
{
    const Int_t nx = m.GetNrows();
    const Int_t ny = m.GetNcols();

    for (Int_t i=0; i<nx; i++)
        for (Int_t j=0; j<ny; j++)
        {
            const Double_t bin =  h.GetBinContent(i+1, j+1);
            m(i, j) = bin*bin;
        }
}

Double_t GetMatrixSumRow(const TMatrixD &m, Int_t row)
{
    const Int_t n = m.GetNcols();

    Double_t sum = 0;
    for (Int_t i=0; i<n; i++)
        sum += m(row, i);

    return sum;
}

Double_t GetMatrixSumDiag(const TMatrixD &m)
{
    const Int_t n = m.GetNcols();

    Double_t sum = 0;
    for (Int_t i=0; i<n; i++)
        sum += m(i, i);

    return sum;
}

Double_t GetMatrixSumCol(const TMatrixD &m, Int_t col=0)
{
    const Int_t n = m.GetNrows();

    Double_t sum = 0;
    for (Int_t i=0; i<n; i++)
        sum += m(i, col);

    return sum;
}
Double_t GetMatrixSum(const TMatrixD &m)
{
    const Int_t n = m.GetNrows();

    Double_t sum = 0;
    for (Int_t i=0; i<n; i++)
        sum += GetMatrixSumRow(m, i);

    return sum;
}


// ======================================================
//
// SteerUnfold
//
void SteerUnfold(TString bintitle,
                 TH1D &ha,     TH2D &hacov, TH2D &hmig,
                 TH2D &hmigor, TH1D &hb0,   TH1D *hpr,
                 TH1D &hb)
{
    // ha      is the distribution to be unfolded
    // hacov   is the covariance matrix of the distribution ha
    // hmig    is the migration matrix;
    //         it is used in the unfolding unless it is overwritten
    //         by SmoothMigrationMatrix by the smoothed migration matrix
    // hmigor  is the migration matrix to be smoothed;
    //         the smoothed migration matrix will be used in the unfolding
    // hpr     the prior distribution
    //         it is only used if SetPriorInput(*hpr) is called   
    // hb      unfolded distribution

    //..............................................       
    // create an MUnfold object;
    // fill histograms into vectors and matrices

    MUnfold unfold(ha, hacov, hmig);
    unfold.bintitle = bintitle;

    //..............................................       
    // smooth the migration matrix;
    //        the smoothed migration matrix will be used in the unfolding
    //        hmig is the original (unsmoothed) migration matrix

    unfold.SmoothMigrationMatrix(hmigor);

    //..............................................       
    // define prior distribution (has always to be defined) 
    // the alternatives are  :

    // 1  SetPriorConstant() :   isotropic distribution
    // 2  SetPriorPower(gamma) : dN/dE = E^{-gamma}
    // 3  SetPriorInput(*hpr):   the distribution *hpr is used 
    // 4  SetPriorRebin(*ha) :   use rebinned histogram ha 

    UInt_t flagprior = 4;
    cout << "SteerUnfold : flagprior = " << flagprior << endl;
    cout << "==========================="<< endl;

    Bool_t errorprior=kTRUE;
    switch (flagprior)
    {
    case 1:
        unfold.SetPriorConstant();
        break;
    case 2:
        errorprior = unfold.SetPriorPower(1.5);
        break;
    case 3:
        if (!hpr)
        {
            cout << "Error: No hpr!" << endl;
            return;
        }
        errorprior = unfold.SetPriorInput(*hpr);
        break;
    case 4:
        errorprior = unfold.SetPriorRebin(ha);
        break;
    }
    if (!errorprior)
    {
        cout << "MUnfold::SetPrior... : failed.  flagprior = " ;
        cout << flagprior << endl;
        return;
    }

    //..............................................       
    // calculate the matrix G = M * M(transposed)
    //           M being the migration matrix

    unfold.CalculateG();

    //..............................................       
    // call steering routine for the actual unfolding;
    // the alternatives are :

    // 1  Schmelling : minimize the function Z by Gauss-Newton iteration;
    //                 the parameters to be fitted are gamma(i) = lambda(i)/w;

    // 2  Tikhonov2 :  regularization term is sum of (2nd deriv.)**2 ;
    //                 minimization by using MINUIT;
    //                 the parameters to be fitted are
    //                 the bin contents of the unfolded distribution

    // 3  Bertero:     minimization by iteration
    //                 

    UInt_t flagunfold = 1;
    cout << "SteerUnfold : flagunfold = "  << flagunfold << endl;
    cout << "============================" << endl;



    switch (flagunfold)
    {
    case 1: // Schmelling
        cout << "" << endl;
        cout << "Unfolding algorithm : Schmelling" << endl;
        cout << "================================" << endl;
        if (!unfold.Schmelling(hb0))
            cout << "MUnfold::Schmelling : failed." << endl;
        break;

    case 2: // Tikhonov2
        cout << "" << endl;
        cout << "Unfolding algorithm :   Tikhonov" << endl;
        cout << "================================" << endl;
        if (!unfold.Tikhonov2(hb0))
            cout << "MUnfold::Tikhonov2 : failed." << endl;
        break;

    case 3: // Bertero
        cout << "" << endl;
        cout << "Unfolding algorithm :    Bertero" << endl;
        cout << "================================" << endl;
        if (!unfold.Bertero(hb0))
            cout << "MUnfold::Bertero : failed." << endl;
        break;
    }    


    //..............................................       
    // Print fResult
    unfold.PrintResults();


    //..............................................       
    // Draw the plots
    unfold.DrawPlots();

    //..............................................       
    // get unfolded distribution
    TMatrixD &Vb    = unfold.GetVb();
    TMatrixD &Vbcov = unfold.GetVbcov();

    UInt_t fNb = unfold.fNb;

    for (UInt_t a=0; a<fNb; a++)
    {
      hb.SetBinContent(a+1, Vb(a,0));
      hb.SetBinError(a+1, sqrt(Vbcov(a, a)) );
    }

}

//__________________________________________________________________________
////////////////////////////////////////////////////////////////////////////
//                                                                        //
// doUnfolding    (to be called in the analysis)                          //
//                                                                        //
// arguments :                                                            //
//                                                                        //
//   INPUT                                                                //
//      TH2D &tobeunfolded : no.of excess events and its error            //
//                           vs. (E-est, Theta)                           //
//      TH3D &migration    : migration matrix (E-est, E_true, Theta)      //
//                                                                        //
//   OUITPUT                                                              //
//      TH2D &unfolded     : no.of excess events and its error            //
//                           vs. (E-true, Theta)                          //
//                                                                        //
// calls SteerUnfold to do the unfolding                                  //
//                                                                        //
// The "Theta" axis is only used to loop over the bins of theta           //
// and to do the unfolding for each bin of theta. Instead of theta        //
// any other variable (or a dummy variable) may be used.                  //
//                                                                        //
//                                                                        //
////////////////////////////////////////////////////////////////////////////

void doUnfolding(TH2D &tobeunfolded, TH3D &migration, TH2D &unfolded)
{
  TAxis &taxis  = *tobeunfolded.GetYaxis();
  Int_t numybins = taxis.GetNbins();

  for (Int_t m=1; m<=numybins; m++)
  {
    TString bintitle = "Bin ";
    bintitle += m;
    bintitle += ": ";

    // -----------------------------------------
    // ha : distribution to be unfolded

    TH1D &ha = *tobeunfolded.ProjectionX("", m, m, "e");
    TString title = bintitle;
    title += "E-est distr. to be unfolded";
    ha.SetNameTitle("ha", title);
    TAxis &aaxis = *ha.GetXaxis();
    Int_t na = aaxis.GetNbins();
    Double_t alow = aaxis.GetBinLowEdge(1);
    Double_t aup  = aaxis.GetBinLowEdge(na+1);

    PrintTH1Content(ha);
    PrintTH1Error(ha);

    // -----------------------------------------
    // covariance matrix of the distribution ha

    title = bintitle;
    title +=  "Error matrix of distribution ha";
    TH2D hacov("hacov", title, na, alow, aup, na, alow, aup);
    //MH::SetBinning(&hacov, &aaxis, &aaxis); 

    Double_t errmin = 3.0;
    for (Int_t i=1; i<=na; i++)
    {
      for (Int_t j=1; j<=na; j++)
      {
        hacov.SetBinContent(i, j, 0.0);
      }
      const Double_t content = ha.GetBinContent(i);
      const Double_t error2  = (ha.GetBinError(i))*(ha.GetBinError(i));
      if (content <= errmin  &&  error2 < errmin) 
        hacov.SetBinContent(i, i, errmin);
      else
        hacov.SetBinContent(i, i, error2);
    }

    //PrintTH2Content(hacov);
    

    // -----------------------------------------
    // migration matrix :
    //           x corresponds to measured quantity
    //           y corresponds to true quantity
    TH2D &hmig = *(TH2D*)migration.Project3D("yxe");
    title = bintitle;
    title += "Migration Matrix";
    hmig.SetNameTitle("Migrat", title);

    TAxis &aaxismig = *hmig.GetXaxis();
    Int_t namig = aaxismig.GetNbins();

    if (na != namig)
    {
      cout << "doUnfolding : binnings are incompatible; na, namig = "
           << na << ",  " << namig << endl;
      return;
    }

    TAxis &baxismig = *hmig.GetYaxis();
    Int_t nbmig = baxismig.GetNbins();
    Double_t blow = baxismig.GetBinLowEdge(1);
    Double_t bup  = baxismig.GetBinLowEdge(nbmig+1);

    PrintTH2Content(hmig);
    //PrintTH2Error(hmig);


    // -----------------------------------------
    // dummy ideal distribution

    Int_t nb = nbmig;

    title = bintitle;
    title += "Dummy Ideal distribution";
    TH1D hb0("dummyhb0", title, nb, blow, bup);;
    //MH::SetBinning(&hb0, &baxismig); 
    hb0.Sumw2();

    for (Int_t k=1; k<=nb; k++)
    {
        hb0.SetBinContent(k, 1.0/nb);
        hb0.SetBinError  (k, 0.1/nb);
    }

    //PrintTH1Content(hb0);

    // -----------------------------------------
    // here the prior distribution can be defined for the call
    // to SetPriorInput(*hpr)

    title = bintitle;
    title += "Dummy Prior distribution";
    TH1D hpr("hpr", title, nb, blow, bup);
    //MH::SetBinning(&hpr, &baxismig); 
    
    for (Int_t k=1; k<=nb; k++)
        hpr.SetBinContent(k, 1.0/nb);

    //PrintTH1Content(hpr);


    // -----------------------------------------
    // unfolded distribution


    title = bintitle;
    title += "Unfolded distribution";
    TH1D hb("hb", title, nb, blow, bup);
    //MH::SetBinning(&hb, &baxismig); 

    // -----------------------------------------
    SteerUnfold(bintitle, ha, hacov, hmig, hmig, hb0, &hpr, hb);

    for (Int_t k=1; k<=nb; k++)
    {
      Double_t content = hb.GetBinContent(k);
      Double_t error   = hb.GetBinError(k);

      unfolded.SetBinContent(k, m, content);
      unfolded.SetBinError(k, m, error);
    }    

    delete &ha;
    delete &hmig;
  }

}
//========================================================================//


////////////////////////////////////////////////////////////////////////////
//                                                                        //
// Main program   (for testing purposes)                                  //
//                                                                        //
//                defines the ideal distribution          (hb0)           //
//                defines the migration matrix            (hMigrat)       //
//                defines the distribution to be unfolded (hVa)           //
//                                                                        //
//                calls doUnfolding                                       //
//                      to do the unfolding                               //
//                                                                        //
////////////////////////////////////////////////////////////////////////////
void fluxMUnfold()
{
  // -----------------------------------------
  // migration matrix :
  //           x corresponds to measured quantity
  //           y corresponds to true quantity

    const Int_t  na   =   13;
    //const Int_t  na   =   18;
    const Axis_t alow = 0.25;
    const Axis_t aup  = 3.50;

    const Int_t  nb   =   11;
    //const Int_t  nb   =   22;
    const Axis_t blow = 0.50;
    const Axis_t bup  = 3.25;

    const Int_t  nc   =     1;
    const Axis_t clow =   0.0;
    const Axis_t cup  =   1.0;

    Int_t m = 1;

    TH3D migration("migration", "Migration Matrix", 
                   na, alow, aup, nb, blow, bup, nc, clow, cup);
    migration.Sumw2();

    // parametrize migration matrix as
    //     <log10(Eest)>      = a0 + a1*log10(Etrue) + a2*log10(Etrue)**2 
    //                             + log10(Etrue) 
    //     RMS( log10(Eest) ) = b0 + b1*log10(Etrue) + b2*log10(Etrue)**2
    Double_t a0 = 0.0;
    Double_t a1 = 0.0;
    Double_t a2 = 0.0;

    Double_t b0 = 0.26;
    Double_t b1 =-0.054;
    Double_t b2 = 0.0;

    TF1 f2("f2", "gaus(0)", alow, aup);
    f2.SetParName(0, "ampl");
    f2.SetParName(1, "mean");
    f2.SetParName(2, "sigma");

    // loop over log10(Etrue) bins
    TAxis &yaxis = *migration.GetYaxis();
    for (Int_t j=1; j<=nb; j++)
    {
        Double_t yvalue = yaxis.GetBinCenter(j);

        const Double_t mean  = a0 + a1*yvalue + a2*yvalue*yvalue + yvalue;
        const Double_t sigma = b0 + b1*yvalue + b2*yvalue*yvalue;
        const Double_t ampl  = 1./ ( sigma*TMath::Sqrt(2.0*TMath::Pi()));

        // gaus(0) is a substitute for [0]*exp( -0.5*( (x-[1])/[2] )**2 )
        f2.SetParameter(0, ampl);
        f2.SetParameter(1, mean);
        f2.SetParameter(2, sigma);

        // fill temporary 1-dim histogram with the function
        // fill the histogram using
        //    - either FillRandom
        //    - or using Freq
        TH1D htemp("temp", "temp", na, alow, aup);
        htemp.Sumw2();

        for (Int_t k=0; k<1000000; k++)
            htemp.Fill(f2.GetRandom());

        // copy it into the migration matrix
        Double_t sum = 0;
        for (Int_t i=1; i<=na; i++)
        {
            const Stat_t content = htemp.GetBinContent(i);
            migration.SetBinContent(i, j, m, content);
            sum += content;
        }

        // normalize migration matrix
        if (sum==0)
            continue;

        for (Int_t i=1; i<=na; i++)
        {
            const Stat_t content = migration.GetBinContent(i,j,m);
            migration.SetBinContent(i,j,m,       content/sum);
            migration.SetBinError  (i,j,m, sqrt(content)/sum);
        }
    }

    //PrintTH3Content(migration);
    //PrintTH3Error(migration);

    // -----------------------------------------
    // ideal distribution

    TH1D hb0("hb0", "Ideal distribution", nb, blow, bup);
    hb0.Sumw2();

    // fill histogram with random numbers according to 
    // an exponential function dN/dE = E^{-gamma}
    //    or with y = log10(E), E = 10^y :             
    //                          dN/dy = ln10 * 10^{y*(1-gamma)}     
    TF1 f1("f1", "pow(10.0, x*(1.0-[0]))", blow, bup);
    f1.SetParName(0,"gamma");
    f1.SetParameter(0, 2.7);

    // ntimes is the number of entries
    for (Int_t k=0; k<10000; k++)
        hb0.Fill(f1.GetRandom());

    // introduce energy threshold at 50 GeV

    const Double_t  lgEth = 1.70;
    const Double_t dlgEth = 0.09;

    for (Int_t j=1; j<=nb; j++)
    {
        const Double_t lgE = hb0.GetBinCenter(j);
        const Double_t c   = hb0.GetBinContent(j);
        const Double_t dc  = hb0.GetBinError(j);
        const Double_t f   = 1.0 / (1.0 + exp( -(lgE-lgEth)/dlgEth ));

        hb0.SetBinContent(j, f* c);
        hb0.SetBinError  (j, f*dc);
    }

    //PrintTH1Content(hb0);

    // -----------------------------------------
    // generate distribution to be unfolded (ha)
    // by smearing the ideal distribution  (hb0)
    //
    TH2D tobeunfolded("tobeunfolded", "Distribution to be unfolded", 
                      na, alow, aup, nc, clow, cup);
    tobeunfolded.Sumw2();

    for (Int_t i=1; i<=na; i++)
    {
        Double_t cont = 0;
        for (Int_t j=1; j<=nb; j++)
            cont += migration.GetBinContent(i, j, m) * hb0.GetBinContent(j);

        tobeunfolded.SetBinContent(i, m, cont);
        tobeunfolded.SetBinError(i, m, sqrt(cont));
    }

    //PrintTH2Content(tobeunfolded);
    //PrintTH2Error(tobeunfolded);

    // -----------------------------------------
    // unfolded distribution

    TH2D unfolded("unfolded", "Unfolded distribution", 
                  nb, blow, bup, nc, clow, cup);
    unfolded.Sumw2();

    // -----------------------------------------
    doUnfolding(tobeunfolded, migration, unfolded);

}
//========================================================================//