#include <iostream>
#include <fstream>
#include <stdio.h>
#include <math.h>
#include <TH1F.h>
#include <TH2F.h>
#include <TH1D.h>
#include <TH2D.h>
#include <TFile.h>
#include <TROOT.h>
#include <TList.h>
#include <TString.h>
//#include <TObjectString.h>
#include <TGraphAsymmErrors.h>
#include <TGraphErrors.h>
#include <TChain.h>
#include <TCutG.h>
#include <TF1.h>
#include <TCanvas.h>
#include <TObjString.h>
#include <TMath.h>

#include <PamUnfold.h>

using namespace std;

ClassImp(PamUnfold);


PamUnfold::PamUnfold(TString name, TString title) : TNamed(name, title){

  cout << "WARNING:: entering PamUnfold::PamUnfold(TString name, TString title)" << endl;
  cout << "           empty constructor be sure to initialize measured and smearing" << endl;

  _measured = NULL;
  _smearing = NULL;
  _prior = NULL;

}

PamUnfold::~PamUnfold(){
}

PamUnfold::PamUnfold(TString name, TString title, TH1D* measured, TH2D* smearing) : TNamed(name, title){

  _measured = measured;
  _smearing = smearing;
  _prior = NULL;

  if( !IsBinningOK() )
    cerr << " -- ERROR in PamUnfold::PamUnfold(TString name, TString title, TH1D* measured, TH2D* smearing)" << endl;

  if( !IsSmNormalized() ){
    cerr << " -- WARNING in PamUnfold::PamUnfold(TString name, TString title, TH1D* measured, TH2D* smearing)" << endl;
    cerr << " ---- Remember to provide the normalization histogram for the smearing matrix" << endl;
  }

  Init();

  _is_improved = kFALSE;
  _smooth = kFALSE;
  _smooth_opt = "ROOT";

  _nsamples = 500;
  _max_steps = 50;
  _min_chi2 = 1.;

}

void PamUnfold::AddExcludedBin(Int_t bin){
  _excluded_bins.push_back(bin);
}

TH1D* PamUnfold::GetMeasured(){
  return _measured;
}

TH2D* PamUnfold::GetSmearing(){
  return _smearing;
}

TH1D* PamUnfold::GetUnfolded(){
  TH1D* __unfolded = (TH1D*) _unfolded->Clone( Form("%s_%s_unf", _measured->GetName(), this->GetName()) );
  return __unfolded;
}

TList* PamUnfold::GetBinHistList(){
  return _bin_hist_list;
}

void PamUnfold::SetMeasured(TH1D* measured){
  _measured = measured;
}

void PamUnfold::SetSmearing(TH2D* smearing){

  _smearing = smearing;

  if( !IsBinningOK() )
    cerr << " -- ERROR in PamUnfold::SetSmearing(TH2D* smearing)" << endl;

  Init();

}

void PamUnfold::SetPrior(TH1D* prior){

  _prior = (TH1D*) prior->Clone( Form("_prior_%s", this->GetName()) );

  if( !IsBinningOK() )
    cerr << " -- ERROR in PamUnfold::SetPrior(TH1D* prior)" << endl;

}

void PamUnfold::SetNormalization(TH1D* norm){
  _norm = norm;

  //If the matrix is not normalized it is normalized now.
  if( !IsSmNormalized() ){
    cout << "      Normalizing smearing matrix" << endl;

    try{ if(!_norm) throw 1; }
    catch(int e){
      if(e==1){
	cerr << " -- ERROR in PamUnfold::Init()" << endl;
	cerr << " ---- Normalization histogram is NULL" << endl;
	cerr << "      execution is likely to crash." << endl;
      }
    }

    NormalizeMatrix();

  }

}

void PamUnfold::SetImproved(Bool_t improv){
  _is_improved = improv;
}

void PamUnfold::SetSmoothing(Bool_t smooth, TString smooth_opt){

  _smooth = smooth;
  _smooth_opt = smooth_opt;

  if(!_smooth_opt.CompareTo(""))
    _smooth_opt = "ROOT";

}

void PamUnfold::SetNsamples(UInt_t nsamples){
  _nsamples = nsamples;
}

void PamUnfold::SetMaxSteps(UInt_t max_steps){
  _max_steps = max_steps;
}

void PamUnfold::SetMinChi2(Double_t min_chi2){
  _min_chi2 = min_chi2;
}


Bool_t PamUnfold::IsSmNormalized(){

  Double_t sum = 0;
  for(UInt_t i=0; i<_smearing->GetNbinsX(); i++){
    sum = 0;
    for(UInt_t j=0; j<_smearing->GetNbinsY(); j++)
      sum += _smearing->GetBinContent(_smearing->GetBin(i+1,j+1));
    if(sum>1.0000001){
      cout << "Bin: " << i+1 << " sum=" << sum << endl;
      return kFALSE;
    }
  }

  return kTRUE;

}

void PamUnfold::WMovAvSmooth(TH1D* input, vector<Int_t>&excl){

  Double_t* xarr = new Double_t [input->GetNbinsX()];

  for(UInt_t ib=0; ib<input->GetNbinsX(); ib++)
    xarr[ib] = (Double_t) input->GetBinContent(ib+1);

  Bool_t excluding;
  for(UInt_t ib=1; ib<input->GetNbinsX()-1; ib++){

    excluding = kFALSE;
    for(UInt_t iex=0; iex<excl.size(); iex++)
      if(
	 ib   == excl[iex] ||
	 ib+1 == excl[iex] ||
	 ib+2 == excl[iex]
	 ) excluding = kTRUE;
    if(excluding) continue;

    Double_t xc = xarr[ib];
    Double_t xl = xarr[ib-1];
    Double_t xh = xarr[ib+1];
    Double_t x = 0.25*(xh+xl) + 0.5*xc;
    input->SetBinContent(ib+1, (Float_t) x);
  }

  delete xarr;

  return;

}

Double_t PamUnfold::GetChi2H( TH1D* h1, TH1D* h2 ){

  if(h1->GetNbinsX() != h2->GetNbinsX()){
    cout << " -- Warning in PamUnfold::GetChi2H(TH1D*, TH1D*) : Histograms have different number of bins" << endl;
    return -1;
  }

  const Double_t* x_1 = h1->GetXaxis()->GetXbins()->GetArray();
  const Double_t* x_2 = h2->GetXaxis()->GetXbins()->GetArray();

  for(Int_t ib=0; ib < h1->GetNbinsX()+1; ib++){
    if(x_1[ib] != x_2[ib]){
      cout << " -- Warning in PamUnfold::GetChi2H(TH1D*, TH1D*) : Histograms have different number of bins" << endl;
      return -1;
    }
    else
      continue;
  }



  Double_t chi2 = 0;
  for(UInt_t i=0; i<h1->GetNbinsX(); i++){
    if(h1->GetBinContent(i+1) + h2->GetBinContent(i+2) > 1)
      chi2 += pow(h1->GetBinContent(i+1) - h2->GetBinContent(i+1),2)/(h1->GetBinContent(i+1) + h2->GetBinContent(i+2));
    else
      chi2 += pow(h1->GetBinContent(i+1) - h2->GetBinContent(i+1),2);
  }

  return chi2;
}

void PamUnfold::Unfold(){


  //Smoothing is applied to the spectrum, not to the counts histogram!!!
  // ------------------------------------------------------
  //   Prior initialization
  // ------------------------------------------------------

  for(Int_t i=0; i<_prior->GetNbinsX(); i++) 
    _prior->SetBinContent(i+1, _prior->GetBinContent(i+1)/_prior->GetBinWidth(i+1) );

  if(_smooth){
    if(_smooth_opt.Contains("WMA")){
      cout << " --- Using Weighted Moving Average smoothing\n";
      WMovAvSmooth(_prior, _excluded_bins);
    }
    else if(_smooth_opt.Contains("ROOT")){
      cout << " --- Using standard ROOT smoothing\n";
      _prior->Smooth();
    }
    else
      cout << " --- WARNING: No valid smoothing option specified" << endl;
  }

  for(Int_t i=0; i<_prior->GetNbinsX(); i++) 
    _prior->SetBinContent(i+1, _prior->GetBinContent(i+1)*_prior->GetBinWidth(i+1) );

  _prior->Scale(1./_prior->GetSumOfWeights()); //The prior is a 'probability' so it has to be normalized at the very last step

  // ------------------------------------------------------
  //   Unfolding
  // ------------------------------------------------------

  cout << " --- UNFOLDING!!" << endl;

  TH2D* theta = (TH2D*) _smearing->Clone("theta");
  theta->Reset();

  for(Int_t i=0; i<theta->GetNbinsY(); i++){
    for(Int_t j=0; j<theta->GetNbinsX(); j++){

      Double_t s = _smearing->GetBinContent(_smearing->GetBin(i+1,j+1))*_prior->GetBinContent(i+1);
      Double_t ls = 0;
      for(Int_t k=0; k<_smearing->GetNbinsX(); k++)
        ls+=_smearing->GetBinContent(_smearing->GetBin(k+1,j+1))*_prior->GetBinContent(k+1);

      if(ls)
	theta->SetBinContent(theta->GetBin(j+1,i+1),s/ls);
    }
  }

  _unfolded->Reset();

  for(Int_t i=0; i<_unfolded->GetNbinsX(); i++){
    Double_t num = 0;
    Double_t den = 0;
    Double_t err = 0;
    for(Int_t j=0; j<theta->GetNbinsX(); j++){

      num+=theta->GetBinContent( theta->GetBin(j+1,i+1) )*_measured->GetBinContent(j+1);
      den+=_smearing->GetBinContent(_smearing->GetBin(i+1,j+1));
      err+=pow( theta->GetBinContent(theta->GetBin(j+1,i+1))*_measured->GetBinError(j+1), 1 );
      //      err+=pow( theta->GetBinContent(theta->GetBin(j+1,i+1))*sp->GetBinError(j+1), 2 );

    }

    if(den){
      _unfolded->SetBinContent(i+1, num/den);
      _unfolded->SetBinError(i+1, fabs(err)/den);
      //      _unfolded->SetBinError(i+1,sqrt(err)/den);
    }
  }

}


void PamUnfold::ImprovedUnfold(){


  vector<Int_t> mu(_measured->GetNbinsX());

  //Smoothing is applied to the spectrum, not to the counts histogram!!!
  // ------------------------------------------------------
  //   Prior initialization
  // ------------------------------------------------------

  for(Int_t i=0; i<_prior->GetNbinsX(); i++) 
    _prior->SetBinContent(i+1, _prior->GetBinContent(i+1)/_prior->GetBinWidth(i+1) );

  if(_smooth){
    if(_smooth_opt.Contains("WMA")){
      cout << " --- Using Weighted Moving Average smoothing\n";
      WMovAvSmooth(_prior, _excluded_bins);
    }
    else if(_smooth_opt.Contains("ROOT")){
      cout << " --- Using standard ROOT smoothing\n";
      _prior->Smooth();
    }
    else
      cout << " --- WARNING: No valid smoothing option specified" << endl;
  }

  for(Int_t i=0; i<_prior->GetNbinsX(); i++) 
    _prior->SetBinContent(i+1, _prior->GetBinContent(i+1)*_prior->GetBinWidth(i+1) );

  _prior->Scale(1./_prior->GetSumOfWeights()); //The prior is a 'probability' so it has to be normalized at the very last step

  // ------------------------------------------------------
  //   Unfolding
  // ------------------------------------------------------

  cout << " --- UNFOLDING!!" << endl;

  _bin_list->Clear();

  TMVA::Timer* timer = new TMVA::Timer(_nsamples, "unf");
  for(UInt_t sample=0; sample<_nsamples; sample++){
    
    cout << " ---- Sampling... " << timer->GetLeftTime(sample) << " left..." << endl << flush <<"\33[1A";

    SampleMatrix();
    
    TH2D* theta = (TH2D*) _smearing_sample->Clone("theta");
    theta->Reset();
    
    for(Int_t i=0; i<theta->GetNbinsY(); i++){
      for(Int_t j=0; j<theta->GetNbinsX(); j++){
	
	Double_t s = _smearing_sample->GetBinContent(_smearing_sample->GetBin(i+1,j+1))*_prior->GetBinContent(i+1);
	Double_t ls = 0;
	for(Int_t k=0; k<_smearing_sample->GetNbinsX(); k++)
	  ls+=_smearing_sample->GetBinContent(_smearing_sample->GetBin(k+1,j+1))*_prior->GetBinContent(k+1);
	
	if(ls)
	  theta->SetBinContent(theta->GetBin(j+1,i+1),s/ls);

      }
    }
    
    TVectorD* result = new TVectorD(_unfolded->GetNbinsX());
    
    //Sampling mu_j and rounding to nearest integer
    for(Int_t j=0; j<_measured->GetNbinsX(); j++){
      mu[j] = TMath::Nint( rangen->Gamma( 1 + _measured->GetBinContent(j+1), 1 ) );

      vector<Double_t> theta_vec_j(_unfolded->GetNbinsX());
      for(Int_t i=0; i<_unfolded->GetNbinsX(); i++){
	theta_vec_j[i] = theta->GetBinContent( theta->GetBin(j+1,i+1) );
	//	cout << theta_vec_j[i] << " ";
      }
      //      cout << endl;

      vector<Int_t> res_partial = rangen->Multinomial(mu[j], theta_vec_j);
      for(Int_t i=0; i<_unfolded->GetNbinsX(); i++){
	//	cout << res_partial[i] << " ";
	(*result)[i] += res_partial[i];
      }
      // cout << endl;
      // cout << _measured->GetBinContent(j+1) << " " << mu[j] << " " << res_partial[17] << " " << (*result)[17] << endl;

    }

    
    for(Int_t i=0; i<_unfolded->GetNbinsX(); i++){
      Double_t den = 0;
      for(Int_t j=0; j<_measured->GetNbinsX(); j++)
	den+=_smearing_sample->GetBinContent(_smearing_sample->GetBin(i+1,j+1));   

	(*result)[i] /= den;
    }

    _bin_list->Add(result);
    
  }

  BuildFlux();

}


Bool_t PamUnfold::IsBinningOK(){

  Int_t n_meas = _measured->GetNbinsX();
  Int_t n_smy = _smearing->GetNbinsY();


  if(n_meas != n_smy){
    cerr << " --- Different binning between measured and smearing" << endl;
    cerr << " ---- measured: " << n_meas << " bins;  smearing:" << n_smy << " bins on Y-axis" << endl;
    return kFALSE;
  }

  const Double_t* x_meas = _measured->GetXaxis()->GetXbins()->GetArray();
  const Double_t* y_sm = _smearing->GetYaxis()->GetXbins()->GetArray();

  for(Int_t ib=0; ib < n_meas+1; ib++){
    if( fabs(x_meas[ib] - y_sm[ib])/x_meas[ib] > 1e-4 ){
      cerr << " --- Different binning between measured and smearing" << endl;
      cerr << " ---- x_meas[" << ib << "] = " << x_meas[ib] << ";"
	   <<"  y_sm[" << ib << "] = " << y_sm[ib] << ";" << endl;
      return kFALSE;
    }
    else
      continue;
  }



  if(_prior){

    Int_t n_prior = _prior->GetNbinsX();
    Int_t n_smx = _smearing->GetNbinsX();

    if(n_prior != n_smx){
      cerr << " --- Different binning between prior and smearing" << endl;
      return kFALSE;
    }

    const Double_t* x_prior = _prior->GetXaxis()->GetXbins()->GetArray();
    const Double_t* x_sm = _smearing->GetXaxis()->GetXbins()->GetArray();
    
    for(Int_t ib=0; ib < n_prior+1; ib++){
      if( fabs(x_prior[ib] - x_sm[ib])/x_sm[ib] > 1e-4 ){
	cerr << " --- Different binning between prior and smearing" << endl;
	return kFALSE;
      }
      else
	continue;
    }
  }


  return kTRUE;
  
}

void PamUnfold::IterativeUnfolding(UInt_t niter, TList* list){

  cout << " -- PamUnfold object " << this->GetName() << endl
       << "      Starting unfolding";

  if(!niter)
    cout << " with chi2 convergence check";
  else
    cout << " with " << niter << " iterations";

  if(_is_improved)
    cout << " and improved algorithm (may take a while...)";
  else
    cout << " and standard algorithm";

  cout << endl;

  if(list)
    cout << "      Saving iterations to TList object at " << list << endl;


  UInt_t iiter = 0;
  Double_t chi2 = 1e6;
  Bool_t kFlag = kTRUE;

  while( kFlag ){

    if(iiter > 0)
      SetPrior(_old_unfolded);

    if(_is_improved)
      ImprovedUnfold();
    else
      Unfold();

    if(iiter > 0){
      chi2 = GetChi2H( _unfolded, _old_unfolded );
      cout << "       Chi2 of change from iteration " << iiter-1 << " to " << iiter << " = " << chi2 << endl;
    }

    if(!niter)
      kFlag = chi2 < _min_chi2  ?  kFALSE : kTRUE;
    else
      kFlag = iiter == niter  ?  kFALSE : kTRUE;

    if(iiter >= _max_steps){
      kFlag = kFALSE;
      cerr << "WARNING: Unfolding procedure did not converge in less than " << _max_steps << " steps\n";
    }

    _old_unfolded = GetUnfolded();

    _old_unfolded->SetName( Form("%s_%03i", _old_unfolded->GetName(), iiter) );
    if(list){ 
      list->Add( _old_unfolded );
    }
    iiter++;

  }

  cout << "      Unfolding converged" << endl;

}

void PamUnfold::Init(){

  cout << " -- PamUnfold object " << this->GetName() << endl
       << "      Initializing unfolded histogram" << endl;

  _unfolded = new TH1D( Form("%s_%s_u", _measured->GetName(), this->GetName()), "", _smearing->GetNbinsX(), _smearing->GetXaxis()->GetXbins()->GetArray());
  _unfolded->Reset();

  _prior = (TH1D*) _unfolded->Clone( Form("_prior_%s",this->GetName()) );
  _prior->Reset();

  cout << "      Initializing flat Prior" << endl;
  //Initializing flat prior
  for(UInt_t ibin=0; ibin<_prior->GetNbinsX(); ibin++){
    _prior->SetBinContent(ibin+1, _prior->GetBinWidth(ibin+1)/(_prior->GetBinLowEdge(_prior->GetNbinsX()+1)-_prior->GetBinLowEdge(1)) );
  }

  cout << "      Initializing Random number generator" << endl;
  rangen = new RanGen();

  _bin_list = new TList();
  _bin_hist_list = new TList();
  _smearing_sample = (TH2D*) _smearing->Clone( Form("%s_sample", _smearing->GetName()) );

  cout << "      Initialization done" << endl;

}


void PamUnfold::NormalizeMatrix(){

  for(UInt_t ibiny=0; ibiny<_smearing->GetNbinsY(); ibiny++){
    for(UInt_t ibinx=0; ibinx<_smearing->GetNbinsX(); ibinx++){
      Int_t globalbin = _smearing->GetBin(ibinx+1, ibiny+1);
      if(_norm->GetBinContent(ibinx+1))
	_smearing->SetBinContent( globalbin, _smearing->GetBinContent(globalbin)/_norm->GetBinContent(ibinx+1) );
    }
  }

}

void PamUnfold::SampleMatrix(){

  _smearing_sample->Reset();

  vector<Int_t> alpha(_smearing->GetNbinsY()+1);
  vector<Double_t> sampled;

  for(UInt_t ibinx=0; ibinx<_smearing->GetNbinsX(); ibinx++){

    //    cout << _norm->GetBinContent(ibinx+1) << endl;
    Double_t snorm=0;
    for(UInt_t ibiny=0; ibiny<_smearing->GetNbinsY(); ibiny++){
      Int_t globalbin = _smearing->GetBin(ibinx+1, ibiny+1);
      snorm += _smearing->GetBinContent(globalbin);
      alpha[ibiny] = TMath::Nint( _smearing->GetBinContent(globalbin) * _norm->GetBinContent(ibinx+1) ); 
    }
    alpha[_smearing->GetNbinsY()] = TMath::Nint( (1 - snorm) * _norm->GetBinContent(ibinx+1) );
    //cout << "bin: " << ibinx << "  " << alpha[_smearing->GetNbinsY()] << endl;

    sampled = rangen->Dirichlet(_smearing->GetNbinsY() + 1, alpha);

    for(UInt_t ibiny=0; ibiny<_smearing_sample->GetNbinsY(); ibiny++){
      Int_t globalbin = _smearing_sample->GetBin(ibinx+1, ibiny+1);
      //      cout << sampled[ibiny] << " ";
      _smearing_sample->SetBinContent(globalbin, sampled[ibiny]);
    }
    //    cout << endl;
  }
  //  cout << endl;

}


void PamUnfold::BuildFlux(){

  _unfolded->Reset();
  _bin_hist_list->Delete();

  for(Int_t j=0; j<_unfolded->GetNbinsX(); j++){
    Double_t min=1e9;
    Double_t max=0;
    Double_t check=0;

    for(Int_t i=0; i<_bin_list->GetEntries(); i++){
      TVectorD* vv = (TVectorD*) _bin_list->At(i);
      check = (*vv)(j);
      if(check<min) min=check;
      if(check>max) max=check;
    }

    _bin_hist_list->Add( new TH1D( Form("hh_%03i", j), ";Events", 150, min-20, max+20 ) );
  }

  for(Int_t i=0; i<_bin_list->GetEntries(); i++){
    TVectorD* vv = (TVectorD*) _bin_list->At(i);
    //cout << "Bin " << i << " ";
    for(Int_t j=0; j<vv->GetNoElements(); j++){
      //cout << (*vv)(j) << " ";
      ((TH1D*) _bin_hist_list->At(j))->Fill( (*vv)(j) );
    }
    //    cout << endl;
  }

  Double_t qq = 0.682689492137;
  Double_t yq[2], xq[2];
  xq[0] = (1-qq)/2; xq[1] = xq[0] + qq;


  for(Int_t j=0; j<_unfolded->GetNbinsX(); j++){
    TH1D* mt = ((TH1D*) _bin_hist_list->At(j));
    mt->GetQuantiles(2, yq, xq);

    double mean = 0.5*(yq[0] + yq[1]);
    double err  = 0.5*(yq[1] - yq[0]);
    if(mt->ComputeIntegral()){
      _unfolded->SetBinContent(j+1, mean);
      _unfolded->SetBinError(j+1, err);
    }
    else{
      _unfolded->SetBinContent(j+1, 0);
      _unfolded->SetBinError(j+1, 0);
    }

  }

}

void PamUnfold::Draw(TString path){

  Int_t nx, ny;
  FindSplitting(_unfolded->GetNbinsX() , nx, ny);
  canv = dynamic_cast<TCanvas*>(gDirectory->FindObject( Form("canv_%s", this->GetName()) ));
  if(!canv) canv = new TCanvas(Form("canv_%s", this->GetName()), "Bin Distributions", 0, 0, 4 + 500*nx, 28 + 500*ny);
  canv->Divide(nx, ny);

  for(Int_t ip=0; ip<_unfolded->GetNbinsX(); ip++){
    canv->cd(ip+1);

    TH1D* mt = ((TH1D*) _bin_hist_list->At(ip));
    mt->Draw();

    Double_t mean = _unfolded->GetBinContent(ip+1);
    Double_t err  = _unfolded->GetBinError(ip+1);

    TLine* line1 = new TLine(mean-err, 0, mean-err, mt->GetMaximum());
    TLine* line2 = new TLine(mean+err, 0, mean+err, mt->GetMaximum());
    TLine* line3 = new TLine(mean, 0, mean, mt->GetMaximum());

    line1->SetLineWidth(2);
    line1->SetLineColor(2);
    line2->SetLineWidth(2);
    line2->SetLineColor(2);
    line3->SetLineWidth(2);
    line3->SetLineColor(4);

    line1->Draw("same");
    line2->Draw("same");
    line3->Draw("same");

  }

  if(path) canv->Print(path);

}

void PamUnfold::FindSplitting(Int_t n, Int_t &nx, Int_t &ny){

  Int_t x = TMath::Nint( sqrt(n) );
  if(x*x == n){ nx=ny=x; return; }

  x = floor( sqrt(n) ) + 1;
  for(Int_t y=0; y<=x; y++){
    if(y*x>=n){ ny = y; break;}
  }
  nx = x;

  if(ny>nx){ nx=ny; ny=x;}

}