////////////////////////////////////////////////////////////////////////
//
// guru_test.cpp, written by Gavin Hesketh (ghesketh@fnal.gov), Feb 2007
//
// A c++ interface to the fortran functions of the guru package
// This example code unfolds the standard example
//
// For more information on this code: http://www-d0.fnal.gov/~ghesketh/unfolding/
// GURU algorithm and fortran code written by 
//            Vato Kartvelishvili (v.kartvelishvili@lancaster.ac.uk)
// For more information on GURU, see Arxiv hep-ph/9509307
//
// Change log:
// 13th Feb 2007: original version.
//
////////////////////////////////////////////////////////////////////////



#include <iostream>
#include <iomanip>
#include <math.h>
#include "TFile.h"
#include "TMath.h"
#include "TH2.h"

using namespace std;


// List of all modules not written in C/C++ language... 
extern "C" { 
  void gurmat_(int &nb, float &rbl, float &rbu, int &nx, float &rxl, float &rxu,
	       float &apro, float &atru, float &xini);

  void gurtst_(int &nb, float &rbl, float &rbu, int &nx, float &rxl, float &rxu,
	       float &apro, float &xtst, float &btst, float &btster, float &xi);

  void gurdat_(int &nb, float &btst, float &bdat, float &bdater);

  void gurc2_(int &nx, int &ici, float &c);
  void gurinv_(int &nx, float &c, float &sc, float &vc, int &nep, float &cp);
  
  void gurrot_(int &nb, int &nx, float &atru, float &xini, float &bdat, float &bdater,
	       float &c, float &cp, float &xi, float &atil, float &btil, float &uort,
	       float &sing, float &vort, float &vreg, float &vrm1, float &xcmdm1, 
	       float &ddat);
  void gurtan_(int &nx, float &ddat, float &aneta, int &neta, float &dav, float &dis);
  
  void gurunf_(int &nx, float &sing, float &ddat, float &vreg, float &tau, int &isin, 
	       float &xi, float &xdat, float &xdatcm);

  void gurchi_(int &nx, float &xtst, float &xdat, float &xcmdm1, float &chi2, float &chi3);

  void guresi_(int &nb, int &nx, float &c, float &apro, float &xdat, float &xini, 
	       float &bdat, float &bdater, 
	       float &ddat, float &sing, float &tau, int &isin, float &xi,
	       float &resib, float &resic, float &resix, 
	       float &resid1, float &resid2, float &resid3, float &xef);
	
  void guruni_(int &nb, int &nx, float &atru, float &bdater, float &cp, float &uort, 
	       float &sing, float &vort, float &tau,
	       float &ainv,float & unitb, float &unitx);
  
  
}


int main() {


  /*
C-----------------------------------------------------------------------  
C in apro(i,k),  i must be MEASURED,     k must be  TRUE (GENERATED)
C    atru(i,k): {b, right axis in PAW}     {x, left axis in PAW}
C-----------------------------------------------------------------------  
  */


  TFile * out_file = new TFile("output_example.root", "RECREATE");

    
  float rbl=0.0;
  float rbu=2.0;

  int nx = 400;
  int nb = nx;
  float ax= (float)nx;
  float rxl=0.0;
  float rxu=2.0;


  float apro[nb][nx];
  float atru[nb][nx];
  float xini[nx];


  cout <<"Setting up Blobel example"<<endl;

  //--------------MAKE RESPONSE MATRICES AND INITIAL X VECTOR-----------

  cout <<"Making responce matrix, initial x vector"<<endl;

  gurmat_(nb, rbl, rbu, nx, rxl,  rxu,
	  apro[0][0], atru[0][0], xini[0]);
 
  /*
  cout <<nb<<" "<< rbl<<" "<<rbu<<" "<< nx<<" "<< rxl<<" "<<  rxu<<endl;
 
   
  for(int i=0; i<nx; i++){
     for(int j=0; j<nx; j++)
       cout <<apro[i][j]<<" ";
     cout <<endl;
   }
   for(int i=0; i<nx; i++)
     cout <<xini[i]<<" ";
   cout <<endl;
  */


  //fortran switches rows for columns, so need to fill plots appropriately:
  //  TH2 *prob_mat_plot = new TH2F("01_Prob_Mat_APRO","Prob.Mat. APRO", nb,rbl,rbu, nx,rxl, rxu);
  TH2 *prob_mat_plot = new TH2F("01_Prob_Mat_APRO","Prob.Mat. APRO", nx,rxl, rxu, nb,rbl,rbu);
  
  for(int i=0; i<nx; i++)
    for(int j=0; j<nb; j++)
      prob_mat_plot->SetBinContent(i+1,j+1, apro[j][i]);

  /*
  cout<<endl<<"**********************"<<endl<<endl;
    for(int i=0; i<nx; i++){
      for(int j=0; j<nb; j++)
	cout << prob_mat_plot->GetBinContent(j+1,i+1)<<" ";
      cout <<endl;
    }
  */

  
  //  TH2 *true_mat_plot = new TH2F("02_Tru_Mat_APRO","True.Mat. APRO", nb,rbl,rbu, nx,rxl, rxu);
  TH2 *true_mat_plot = new TH2F("02_Tru_Mat_APRO","True.Mat. APRO", nx,rxl, rxu, nb,rbl,rbu);
  for(int i=0; i<nx; i++)
    for(int j=0; j<nb; j++)
      true_mat_plot->SetBinContent(i+1,j+1, atru[j][i]);
  
  TH1 *xini_plot = new TH1F("03_xini","xini", nx,rxl, rxu);
  for(int j=0; j<nx; j++)
    xini_plot->SetBinContent(j+1, xini[j]);

  out_file->cd();
  prob_mat_plot->Write();
  true_mat_plot->Write();
  xini_plot->Write();



  //------------------BUILD TEST DISTRIBUTION ----------  

  cout <<"Building test distribution"<<endl;
 
 
  float xtst[nx], xtster[nx], xi[nx];
  float btst[nb], btster[nb];


  gurtst_(nb,rbl,rbu,nx,rxl,rxu,apro[0][0],xtst[0],btst[0],btster[0],xi[0]);
  
  for(int i=0; i<nx; i++)
    xtster[i] = sqrt(xtst[i]);
      
  TH1 *xtst_plot = new TH1F("04_xtst","xtst", nx,rxl, rxu);
  for(int j=0; j<nx; j++){
    xtst_plot->SetBinContent(j+1, xtst[j]);
    xtst_plot->SetBinError(j+1, xtster[j]);
  }

  TH1 *btst_plot = new TH1F("05_btst","btst", nx,rxl, rxu);
  for(int j=0; j<nb; j++){
    btst_plot->SetBinContent(j+1, btst[j]);
    btst_plot->SetBinError(j+1, btster[j]);
  }


  out_file->cd();
  xtst_plot->Write();
  btst_plot->Write();




  //------------------BUILD DATA DISTRIBUTION ----------  

  cout <<"Building data distribution"<<endl;
  float bdat[nb],bdater[nb];

  gurdat_(nb,btst[0],bdat[0],bdater[0]);


  TH1 *bdata_plot = new TH1F("06_bdata","bdata", nb,rbl, rbu);
  for(int j=0; j<nb; j++){
    bdata_plot->SetBinContent(j+1, bdat[j]);
    bdata_plot->SetBinError(j+1, bdater[j]);
  }

  out_file->cd();
  bdata_plot->Write();


  cout <<"Finished setting up example"<<endl;

  // return 0;

  //-----------------------------------------------------------------------------
  //------ Now do the unfolding:
  //-----------------------------------------------------------------------------
  
  //for making derivative & inverse:
  int ici;
  float c[nx][nx];
  float cp[nx][nx];
  float sc[nx];
  float vc[nx][nx]; 
  int nep;
  
  //for rotating, rescaling:
  float btil[nb]; 
  float atil[nb][nx];                   // atil=(bdater)^-1*apro
  float vort[nx][nx];                   // orthogonal matrix v
  float uort[nb][nx];                   // orthogonal matrix u
  float xcmdm1[nx][nx];                  // inverse cov. matrix
  float sing[nx];                      // vector of singular values
  float vreg[nx][nx];
  float vrm1[nx][nx];
  float ddat[nx];
  float dabs[nx];

  //esimating optimal tau:
  float aneta[nx];
  int neta;
  float dav;
  float dis;
  
  //tau
  int ni=1;
  int nk=ni*nx;
  
  double tat[nk];

  //unfolding:
  int isin=0;
  float xdat[nx];
  float  xdatcm[nx][nx];                 // covariance matrix	
  
  //calculating chi^2:
  float chi2, chi3;
  float  xisq[nk], xisl[nk];

  //residual vector lengths:
  float resib[nb],resic[nx],resix[nx];     // partial contributions
  float resid1,resid2,resid3, xef;  
  float deriv1[nk], deriv2[nk], deriv3[nk], xeff[nk];
  float probab[nk], proful[nk];
  float resida[nk],residc[nk],residd[nk];  // tau-dependent sums


  for(int iter=1; iter<=1; iter++) {
    
    //---------------Make derivative matrix C and its inverse CP--------
    cout <<"Make derivative matrix C and its inverse CP"<<endl;
    
    if(iter == 1) ici=2;  // second der. mat 
    //      if(iter ==2)ici=0;  // length for 2nd.

    gurc2_(nx,ici,c[0][0]);
    gurinv_(nx,c[0][0],sc[0],vc[0][0],nep,cp[0][0]);
    // inversion changes the input matrix, so one has to fill c once again
    gurc2_(nx,ici,c[0][0]);
    
    cout <<"While inverting C, "<<nep<<" s.v. were used, "<< nx-nep<<" s.v. were discarded."<<endl;
    //if nep is less than nx the matrix is singular and the inversion does
    // not work properly. Unfolding may go on but results will be unsatisfactory.
    
    
    TH2 *derv_c_plot = new TH2F("07_derv_mat_c","Derivative matrix C", nx, 0., ax, nx,0., ax);
    
    for(int i=0; i<nx; i++)
      for(int j=0; j<nx; j++)
	derv_c_plot->SetBinContent(i+1,j+1, c[j][i]);
    
    
    TH2 *derv_c_inv_plot = new TH2F("08_derv_mat_inv_c","Inverse of matrix C", nx, 0., ax, nx,0., ax);
    
    for(int i=0; i<nx; i++)
      for(int j=0; j<nx; j++)
	derv_c_inv_plot->SetBinContent(i+1,j+1, cp[j][i]);
    
 
    out_file->cd();
    derv_c_plot->Write();
    derv_c_inv_plot->Write();
    
    //-------------------------ROTATE AND RESCALE--------------------
    
    //   cout <<"ROTATE AND RESCALE"<<endl;
    
    gurrot_(nb, nx, atru[0][0], xini[0], bdat[0], bdater[0], c[0][0], cp[0][0], xi[0],
	    atil[0][0], btil[0], uort[0][0], sing[0], vort[0][0], vreg[0][0], 
	    vrm1[0][0], xcmdm1[0][0], ddat[0]);
    
    /*   
    for(int i=0; i<nx; i++)
      cout<<sing[i]<<endl;
    */

    for(int i=0; i<nx; i++)
      dabs[i]=fabs(ddat[i]);
	
 

    //--------------fill histos which do not depend on the choice of tau-----     
     
    // TH2 *rscl_mat_plot = new TH2F("09_Rscld_Mat_ATIL","Rescaled Matrix ATIL",nb,rbl,rbu, nx,rxl,rxu);
    TH2 *rscl_mat_plot = new TH2F("09_Rscld_Mat_ATIL","Rescaled Matrix ATIL", nx,rxl,rxu, nb,rbl,rbu);
    for(int i=0; i<nx; i++)
      for(int j=0; j<nb; j++)
	rscl_mat_plot->SetBinContent(i+1,j+1,atil[j][i]); 
    
    TH2 *v_plot = new TH2F("10_V", "V", nx,0.,ax, nx,0.,ax);
    for(int i=0; i<nx; i++)
      for(int j=0; j<nx; j++)
	v_plot->SetBinContent(i+1,j+1,vort[j][i]);

    TH1 *btil_plot = new TH1F("11_btil", "BTIL", nb,rbl,rbu);
    for(int i=0; i<nb; i++)
      btil_plot->SetBinContent(i+1,btil[i]);
    
    TH1 *sing_plot = new TH1F("12_sing", "SING", nx,0,ax);
    for(int i=0; i<nx; i++)
      sing_plot->SetBinContent(i+1,sing[i]);
    
    TH2 *inv_cov_mat_plot = new TH2F("13_inv_cov_mat", "Inverse Covariant Matrix",nx,rxl,rxu, nx,rxl,rxu);
    for(int i=0; i<nx; i++)
      for(int j=0; j<nx; j++)
	inv_cov_mat_plot->SetBinContent(i+1,j+1,xcmdm1[j][i]);
    
    TH1 *abs_ddat_plot = new TH1F("14_abs_ddat", "ABS(DDAT)", nx,0,ax);
    for(int i=0; i<nx; i++)
      abs_ddat_plot->SetBinContent(i+1,dabs[i]);
    

    out_file->cd();
      
    rscl_mat_plot->Write();
    v_plot->Write();
    btil_plot->Write();
    sing_plot->Write();
    inv_cov_mat_plot->Write();
    abs_ddat_plot->Write();    







    //------------------------ESTIMATE OPTIMAL TAU----------------------

    gurtan_(nx, ddat[0], aneta[0], neta, dav, dis);

	//  neta is an estimate of the effective rank of the system

    cout <<endl<<"neta = "<<neta<<",  dav = "<<dav<<",  dis = "<<dis<<endl<<endl;


	cout <<setw(3)<<"ns"<<setw(12)<<"tau"<<setw(12)<<"chi2/xtst"<<
	  setw(12)<<"resid"<<setw(12)<<"dof"<<setw(12)<<"proful"<<
	  setw(12)<<"probab"<<endl;
	//    cout <<" ns    tau        chi2/xtst       resid         dof  proful  probab"<<endl;
    



    //------------------ T A U -----------------------
 
	// int nsetmax=nb;      
    int nsetmax=20;      

    float stepi=1.0/float(ni);
    
    //    for(int nset=1; nset<=nsetmax; nset++) {
    //correct for fortran/c++ difference (starting at 1 vs 0):
    for(int nset=0; nset<nsetmax; nset++) {
      for(int iset=0; iset<=ni-1; iset++) {
	
	int kset=nset*ni+iset;
	
	if(kset<0) { 
	  //iset=ni+1;
	  //nset=nsetmax;
	  continue;
	}

	
	float power=float(iset)/float(ni);

	/*
	  if(nset.ge.1)then
	  tat(kset)=(sing(nset)**(1.-power) * sing(nset+1)**(power))**2
	  else
	  tat(kset)=32.**(ni-iset)*sing(1)**2
	  endif
	*/
	
	//again, some fortran vs c++ differences need correction:
	if(nset>=0)
	  tat[kset]=pow( (pow(sing[nset],(1.-power)) * pow(sing[nset+1],power)), 2);
	else
	  tat[kset]=pow(32.,(ni-iset)) * pow(sing[0],2);
	
	
	float tau=tat[kset];
	  
	//	cout <<endl<<kset<<" "<<nset<<" "<<sing[nset]<<" "<<power<<" "<<sing[nset+1]<<" ";
	//	cout <<ni<<" "<<iset<<" "<<sing[0]<<" "<<tau<<endl;
	

	
	//-------------------------UNFOLDING----------------------
	
	gurunf_(nx, sing[0], ddat[0], vreg[0][0], tau, isin, xi[0], xdat[0], xdatcm[0][0]);


	char plot_name[50];
	if(kset<9)sprintf(plot_name,"15_xdat_0%i", kset+1);
	else sprintf(plot_name,"15_xdat_%i", kset);
	TH1 *xdat_plot = new TH1F(plot_name, "xdat",nx,rxl,rxu);
	for(int i=0; i<nx; i++)
	  xdat_plot->SetBinContent(i+1, xdat[i]);

	if(kset<9)sprintf(plot_name,"16_Cov_Mat_0%i", kset+1);
	else sprintf(plot_name,"16_Cov_Mat_%i", kset);
	TH2 *cov_mat_plot = new TH2F(plot_name, "Covariance Matrix", nx,rxl,rxu, nx,rxl,rxu); 
	for(int i=0; i<nx; i++)
	  for(int j=0; j<nx; j++)
	    cov_mat_plot->SetBinContent(i+1,j+1,xdatcm[j][i]);
	

	out_file->cd();
	xdat_plot->Write();
	cov_mat_plot->Write();

	//---------------------Calculating Chi2's--------------------

	gurchi_(nx, xtst[0], xdat[0], xcmdm1[0][0], chi2, chi3);
	
	if(chi2>1e6)chi2=999999.0;
	xisq[kset]=chi2;             // chi2 of xtst vs xdat with errors of xdat
	xisl[kset]=chi3;             // chi2 of xdat vs xtst with errors of xtst
	
	
	
	
	//------------------Residue vector lengths etc-------------------
	
	guresi_(nb, nx, c[0][0], apro[0][0], xdat[0], xini[0], bdat[0], bdater[0], 
		ddat[0], sing[0], tau, isin, xi[0],
		resib[0], resic[0], resix[0], resid1, resid2, resid3, xef);
	

	//	for(int i=0; i<nx; i++)
	//	  cout <<xi[i]<<" ";
	//	cout <<endl;

	xeff[kset]=xef;
	
	//------tail probability of chi2-distribution: use incomplete Gamma function--


	float arg1 = (float(nx)-xef)/2.;
	float arg2 = (resid3)/2.;
      	probab[kset]=(1-TMath::Gamma(arg1,arg2));

	arg1 = (float(nb)-xef)/2.;
	arg2=(resid1+resid2)/2.;
      	proful[kset]=(1-TMath::Gamma(arg1,arg2));
	
	float dt=(tat[kset]-tat[kset-1])/tat[kset];
	deriv1[kset]=(probab[kset]-probab[kset-1])/dt;
	deriv2[kset]=(deriv1[kset]-deriv1[kset-1])/dt;
	deriv3[kset]=(deriv2[kset]-deriv2[kset-1])/dt;
	
	cout <<setw(3)<< kset+1<<
	  setw(12)<<tau<<
	  setw(12)<<chi2<<
	  setw(12)<<resid3<<
	  setw(12)<<xef<<
	  setw(12)<<proful[kset]<<
	  setw(12)<<probab[kset]<<endl;
	
	resida[kset]=resid1;     //sum of residues^2 for Ax-b
	residc[kset]=resid2;     //sum of residues^2 for C^T*C
	residd[kset]=resid3;     //sum of residues for diagoalized system
	

	if(kset<9) sprintf(plot_name,"17_resib_0%i", kset+1);
	else  sprintf(plot_name,"17_resib_%i", kset);
	TH1* resib_plot = new TH1F(plot_name,"RESIB", nb,rbl,rbu);
	for(int i=0; i<nb; i++)
	  resib_plot->SetBinContent(i+1, resib[i]);

	if(kset<9) sprintf(plot_name,"18_resic_0%i", kset+1);
	else sprintf(plot_name,"18_resic_%i", kset);
	TH1* resic_plot = new TH1F(plot_name,"RESIC", nx,0.,ax);
	for(int i=0; i<nb; i++)
	  resic_plot->SetBinContent(i+1, resic[i]);

	if(kset<9) sprintf(plot_name,"19_resix_0%i", kset+1);
	else sprintf(plot_name,"19_resix_%i", kset);
	TH1* resix_plot = new TH1F(plot_name,"RESIX", nx,0.,ax);
	for(int i=0; i<nb; i++)
	  resix_plot->SetBinContent(i+1, resix[i]);
	
	out_file->cd();
	resib_plot->Write();
	resic_plot->Write();
	resix_plot->Write();

	/*

C---choose a good tau for xi to be used in the second iteration---------

c      if(iter.eq.1)then
c        if(kset.eq.30)then
cC        if(probab(kset).ge.0.1.and.probab(kset-1).lt.0.1) then
c          write(*,*)'xi taken, kset = ', kset
c          do 3300 i=1,nx
c            xi(i)=xdat(i)
c 3300     continue
c        endif
c      endif


	*/

      }//iset loop
    } //nset loop

  }//iter loop
 







  //--Calculate the Pseudoinverse of ATRU for current TAU=(sing(nsetmax))**2 

  //  float ainv[nx][nb], unitb[nb][nb], unitx[nx][nx];
  
  // guruni_(nb, nx, atru[0][0], bdater[0], cp[0][0], uort[0][0], sing[0], vort[0][0], tau,
  //	  ainv, unitb, unitx);
  
  
    
  TH1 *prob_plot = new TH1F("20_prob_dofeff", "Prob(dofeff)",nk,.5,ni*ax+.5);
  for(int i=0; i<nk; i++)
    prob_plot->SetBinContent(i+1, probab[i]);
  
  TH1 *res1_plot = new TH1F("21_residue_1", "Residue 1",nk,.5,ni*ax+.5);
  for(int i=0; i<nk; i++)
    res1_plot->SetBinContent(i+1, resida[i]);
  
  TH1 *res2_plot = new TH1F("22_residue_2", "Residue 2",nk,.5,ni*ax+.5);
  for(int i=0; i<nk; i++)
    res2_plot->SetBinContent(i+1, residc[i]);

  TH1 *res3_plot = new TH1F("23_residue_3", "Residue 3",nk,.5,ni*ax+.5);
  for(int i=0; i<nk; i++)
    res3_plot->SetBinContent(i+1, residd[i]);


  /*
 // TH2 *ainv_plot = new TH2F("24_ainv", "A INV",nx,0.,ax,nb,0.,ab);
  TH2 *ainv_plot = new TH2F("24_ainv", "A INV", nb,0.,ab, nx,0.,ax);
  for(int i=0; i<nb; i++)
    for(int j=0; j<nx; j++)
      ainv_plot->SetBinContent(i+1,j+1, ainv[j][i]);
  
  TH2 *unitb_plot = new TH2F("25_unitb", "UNIT B",nb,0.,ab,nb,0.,ab);
  for(int i=0; i<nb; i++)
    for(int j=0; j<nb; j++)
      unitb_plot->SetBinContent(i+1,j+1, unitb[j][i]);
  
  TH2 *unitx_plot = new TH2F("26_unitx", "UNIT X",nx,0.,ax,nx,0.,ax);
  for(int i=0; i<nx; i++)
    for(int j=0; j<nx; j++)
      unitx_plot->SetBinContent(i+1,j+1, unitx[j][i]);
  */

  
  TH1 *chi2_plot = new TH1F("27_chi2", "Chi2",nk,0.5,ni*ax+0.5);
  for(int i=0; i<nk; i++)
    chi2_plot->SetBinContent(i+1, xisq[i]);
  
  TH1 *chi2_l_plot = new TH1F("28_chi2_l", "Chi2-L",nk,0.5,ni*ax+0.5);
  for(int i=0; i<nk; i++)
    chi2_l_plot->SetBinContent(i+1, xisl[i]);
  
  
  

  out_file->cd();
  
  
  prob_plot->Write();
  res1_plot->Write(); 
  res2_plot->Write();
  res3_plot->Write();
  chi2_plot->Write();
  chi2_l_plot->Write();
  



  return 0;
}