#include "StHbtMaker/CorrFctn/QinvCorrFctnKt_MomAndEnCons.h"
#include <sstream>
#include <cstdio>

#ifdef __ROOT__ 
ClassImp(QinvCorrFctnKt_MomAndEnCons)
#endif

//_________________
QinvCorrFctnKt_MomAndEnCons::QinvCorrFctnKt_MomAndEnCons(char* title, 
							 const int& nbins, 
							 const float& QinvLo, 
							 const float& QinvHi,
							 const int& nCFs, 
							 const float& KtLo, 
							 const float& KtHi,
							 const double& meanPt2, 
							 const double& meanPz2,
							 const double& meanE2, 
							 const double& meanE, 
							 const int& NofParticles) {

  mPartMass = 0.493677;

  mNumberCFs = nCFs;
  mKtMin = KtLo;
  mKtMax = KtHi;
  mDeltaKt = (mKtMax-mKtMin)/mNumberCFs;

  mMeanPt2 = meanPt2;
  mMeanPz2 = meanPz2;
  mMeanE2 = meanE2;
  mMeanE = meanE;
  mNumberOfParticles = NofParticles;
  
  // set up numerator
  //  title = "Num Qinv (MeV/c)";
  stringstream TitNum, TitDen, TitCurrent;
  stringstream TitNumPt1xPt2, TitNumPt1SqrSumPt2Sqr;
  stringstream TitNumPz1xPz2, TitNumPz1SqrSumPz2Sqr;
  stringstream TitNumE1xE2, TitNumE1sumE2;
  stringstream TitNumFull, TitDenFull;

  TitNum.str("");        TitDen.str("");
  TitNumPt1xPt2.str(""); TitNumPt1SqrSumPt2Sqr.str("");
  TitNumPz1xPz2.str(""); TitNumPz1SqrSumPz2Sqr.str("");
  TitNumE1xE2.str("");   TitNumE1sumE2.str("");
  TitNumFull.str("");    TitDenFull.str("");

  TitNum << title << "_Num";
  TitDen << title << "_Den";
  TitNumPt1xPt2 << title << "_TitNumPt1xPt2";
  TitNumPt1SqrSumPt2Sqr << title << "_TitNumPt1SqrSumPt2Sqr";
  TitNumPz1xPz2 << title << "_TitNumPz1xPz2";
  TitNumPz1SqrSumPz2Sqr << title << "_TitNumPz1SqrSumPz2Sqr";
  TitNumE1xE2 << title << "_TitNumE1xE2";
  TitNumE1sumE2 << title << "_TitNumE1sumE2";
  TitNumFull << title << "_TitNumFull";
  TitDenFull << title << "_TitDenFull";

  mNumerator = new StHbt1DHisto[mNumberCFs];
  mDenominator = new StHbt1DHisto[mNumberCFs];
  mNumeratorPt1xPt2 = new StHbt1DHisto[mNumberCFs];
  mNumeratorPt1SqrSumPt2Sqr = new StHbt1DHisto[mNumberCFs];
  mNumeratorPz1xPz2 = new StHbt1DHisto[mNumberCFs];
  mNumeratorPz1SqrSumPz2Sqr = new StHbt1DHisto[mNumberCFs];
  mNumeratorE1xE2 = new StHbt1DHisto[mNumberCFs];
  mNumeratorE1sumE2 = new StHbt1DHisto[mNumberCFs];
  mNumeratorFull = new StHbt1DHisto[mNumberCFs];
  mDenominatorFull = new StHbt1DHisto[mNumberCFs];
  
  for(int i=0; i<mNumberCFs; i++) {
    TitCurrent.str("");
    TitCurrent << TitNum.str() << "_bin_" << i;
    mNumerator[i].SetName(TitCurrent.str().c_str());
    mNumerator[i].SetTitle(TitCurrent.str().c_str());
    mNumerator[i].SetBins(nbins,QinvLo,QinvHi);
    mNumerator[i].SetDirectory(0);
    mNumerator[i].Sumw2();
    
    TitCurrent.str("");
    TitCurrent << TitNumPt1xPt2.str() << "_bin_" << i;
    mNumeratorPt1xPt2[i].SetName(TitCurrent.str().c_str());
    mNumeratorPt1xPt2[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorPt1xPt2[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorPt1xPt2[i].SetDirectory(0);
    mNumeratorPt1xPt2[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitNumPt1SqrSumPt2Sqr.str() << "_bin_" << i;
    mNumeratorPt1SqrSumPt2Sqr[i].SetName(TitCurrent.str().c_str());
    mNumeratorPt1SqrSumPt2Sqr[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorPt1SqrSumPt2Sqr[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorPt1SqrSumPt2Sqr[i].SetDirectory(0);
    mNumeratorPt1SqrSumPt2Sqr[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitNumPz1xPz2.str() << "_bin_" << i;
    mNumeratorPz1xPz2[i].SetName(TitCurrent.str().c_str());
    mNumeratorPz1xPz2[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorPz1xPz2[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorPz1xPz2[i].SetDirectory(0);
    mNumeratorPz1xPz2[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitNumPz1SqrSumPz2Sqr.str() << "_bin_" << i;
    mNumeratorPz1SqrSumPz2Sqr[i].SetName(TitCurrent.str().c_str());
    mNumeratorPz1SqrSumPz2Sqr[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorPz1SqrSumPz2Sqr[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorPz1SqrSumPz2Sqr[i].SetDirectory(0);
    mNumeratorPz1SqrSumPz2Sqr[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitDen.str() << "_bin_" << i;
    mDenominator[i].SetName(TitCurrent.str().c_str());
    mDenominator[i].SetTitle(TitCurrent.str().c_str());
    mDenominator[i].SetBins(nbins,QinvLo,QinvHi);
    mDenominator[i].SetDirectory(0);
    mDenominator[i].Sumw2();
  
    TitCurrent.str("");
    TitCurrent << TitNumE1xE2.str() << "_bin_" << i;
    mNumeratorE1xE2[i].SetName(TitCurrent.str().c_str());
    mNumeratorE1xE2[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorE1xE2[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorE1xE2[i].SetDirectory(0);
    mNumeratorE1xE2[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitNumE1sumE2.str() << "_bin_" << i;
    mNumeratorE1sumE2[i].SetName(TitCurrent.str().c_str());
    mNumeratorE1sumE2[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorE1sumE2[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorE1sumE2[i].SetDirectory(0);
    mNumeratorE1sumE2[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitNumFull.str() << "_bin_" << i;
    mNumeratorFull[i].SetName(TitCurrent.str().c_str());
    mNumeratorFull[i].SetTitle(TitCurrent.str().c_str());
    mNumeratorFull[i].SetBins(nbins,QinvLo,QinvHi);
    mNumeratorFull[i].SetDirectory(0);
    mNumeratorFull[i].Sumw2();

    TitCurrent.str("");
    TitCurrent << TitDenFull.str() << "_bin_" << i;
    mDenominatorFull[i].SetName(TitCurrent.str().c_str());
    mDenominatorFull[i].SetTitle(TitCurrent.str().c_str());
    mDenominatorFull[i].SetBins(nbins,QinvLo,QinvHi);
    mDenominatorFull[i].SetDirectory(0);
    mDenominatorFull[i].Sumw2();
  }
}

//_________________
QinvCorrFctnKt_MomAndEnCons::~QinvCorrFctnKt_MomAndEnCons() {

  delete [] mNumerator;
  delete [] mDenominator;
  delete [] mNumeratorPt1xPt2;
  delete [] mNumeratorPt1SqrSumPt2Sqr;
  delete [] mNumeratorPz1xPz2;
  delete [] mNumeratorPz1SqrSumPz2Sqr;
  delete [] mNumeratorE1xE2;
  delete [] mNumeratorE1sumE2;
  delete [] mNumeratorFull;
  delete [] mDenominatorFull;
}

//_________________
void QinvCorrFctnKt_MomAndEnCons::WriteOutHistos() {

  for(int i=0; i<mNumberCFs; i++) {
    mNumerator[i].Write();
    mDenominator[i].Write();
    mNumeratorPt1xPt2[i].Write();
    mNumeratorPt1SqrSumPt2Sqr[i].Write();
    mNumeratorPz1xPz2[i].Write();
    mNumeratorPz1SqrSumPz2Sqr[i].Write();
    mNumeratorE1xE2[i].Write();
    mNumeratorE1sumE2[i].Write();
    mNumeratorFull[i].Scale(TMath::Power( double(mNumberOfParticles)/double(mNumberOfParticles  - 2.) , 2.));
    mDenominatorFull[i].Scale(TMath::Power( double(mNumberOfParticles)/double(mNumberOfParticles  - 1.) , 4.));

    mNumeratorFull[i].Write();
    mDenominatorFull[i].Write();
  }
}

//_________________
void QinvCorrFctnKt_MomAndEnCons::Finish() {
  /* Ratios might be here */
}

//_________________
StHbtString QinvCorrFctnKt_MomAndEnCons::Report() {

  int mNumeratorEntries = 0;
  int mDenominatorEntries = 0;
  int mNumeratorPt1xPt2Entries = 0;
  int mNumeratorPt1SqrSumPt2SqrEntries = 0;
  int mNumeratorPz1xPz2Entries = 0;
  int mNumeratorPz1SqrSumPz2SqrEntries = 0;
  int mNumeratorE1xE2Entries = 0;
  int mNumeratorE1sumE2Entries = 0;
  int mNumeratorFullEntries = 0;
  int mDenominatorFullEntries = 0;

  for(int i=0; i<mNumberCFs; i++) {
    mNumeratorEntries += (int)mNumerator[i].GetEntries();
    mDenominatorEntries += (int)mDenominator[i].GetEntries();
    mNumeratorPt1xPt2Entries += (int)mNumeratorPt1xPt2[i].GetEntries();
    mNumeratorPt1SqrSumPt2SqrEntries += (int)mNumeratorPt1SqrSumPt2Sqr[i].GetEntries();
    mNumeratorPz1xPz2Entries += (int)mNumeratorPz1xPz2[i].GetEntries();
    mNumeratorPz1SqrSumPz2SqrEntries += (int)mNumeratorPz1SqrSumPz2Sqr[i].GetEntries();
    mNumeratorE1xE2Entries += (int)mNumeratorE1xE2[i].GetEntries();
    mNumeratorE1sumE2Entries += (int)mNumeratorE1sumE2[i].GetEntries();
    mNumeratorFullEntries += (int)mNumeratorFull[i].GetEntries();
    mDenominatorFullEntries += (int)mDenominatorFull[i].GetEntries();
  }


  string stemp ="Qinv Correlation Function with Energy and Momentum Conservation Report:\n";
  char ctemp[100];
  sprintf(ctemp,"Number of entries in numerator:\t%i\n",mNumeratorEntries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in denominator:\t%i\n",mDenominatorEntries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in numeratorPt1xPt2:\t%i\n",mNumeratorPt1xPt2Entries);
  stemp += ctemp; 
  sprintf(ctemp,"Number of entries in numeratorPt1SqrSumPt2Sqr:\t%i\n",mNumeratorPt1SqrSumPt2SqrEntries);
  stemp += ctemp; 
  sprintf(ctemp,"Number of entries in numeratorPz1xPz2:\t%i\n",mNumeratorPz1xPz2Entries);
  stemp += ctemp;   
  sprintf(ctemp,"Number of entries in numeratorPz1SqrSumPz2Sqr:\t%i\n",mNumeratorPz1SqrSumPz2SqrEntries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in numeratorE1xE2:\t%i\n",mNumeratorE1xE2Entries);
  stemp += ctemp; 
  sprintf(ctemp,"Number of entries in numeratorE1sumE2:\t%i\n",mNumeratorE1sumE2Entries);
  stemp += ctemp; 
  sprintf(ctemp,"Number of entries in numeratorFull:\t%i\n",mNumeratorFullEntries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in denominatorFull:\t%i\n",mDenominatorFullEntries);
  stemp += ctemp; 

  StHbtString returnThis = stemp;
  return returnThis;
}

//_________________
void QinvCorrFctnKt_MomAndEnCons::AddRealPair(const StHbtPair* pair) {

  int mIndex = (int)(((fabs(pair->kT())-mKtMin)/mDeltaKt));
  if ( (mIndex>=0) && (mIndex<mNumberCFs) ) {
    double Qinv = fabs(pair->qInv());

    double px1 = pair->track1()->FourMomentum().x();
    double py1 = pair->track1()->FourMomentum().y();
    double pz1 = pair->track1()->FourMomentum().z();
    double E1  = ::sqrt(px1*px1 + py1*py1 + pz1*pz1 + mPartMass*mPartMass);
  
    double px2 = pair->track2()->FourMomentum().x();
    double py2 = pair->track2()->FourMomentum().y();
    double pz2 = pair->track2()->FourMomentum().z();
    double E2  = ::sqrt(px2*px2 + py2*py2 + pz2*pz2 + mPartMass*mPartMass);
  
    double valuePt1xPt2 = (px1*px2 + py1*py2);
    double valuePt1SqrSumPt2Sqr = (px1*px1 + py1*py1 + px2*px2 + py2*py2);

    double valuePz1xPz2 = pz1*pz2;
    double valuePz1SqrSumPz2Sqr = (pz1*pz1 + pz2*pz2);
        
    mNumerator[mIndex].Fill(Qinv);
    mNumeratorPt1xPt2[mIndex].Fill(Qinv,valuePt1xPt2);
    mNumeratorPt1SqrSumPt2Sqr[mIndex].Fill(Qinv,valuePt1SqrSumPt2Sqr);

    mNumeratorPz1xPz2[mIndex].Fill(Qinv,valuePz1xPz2);
    mNumeratorPz1SqrSumPz2Sqr[mIndex].Fill(Qinv,valuePz1SqrSumPz2Sqr);         

    mNumeratorE1xE2[mIndex].Fill(Qinv,E1*E2);
    mNumeratorE1sumE2[mIndex].Fill(Qinv,E1+E2);

    double valueFullNum = exp(-1. * (1./double(mNumberOfParticles-2.) * 
				     ( ((valuePt1SqrSumPt2Sqr+2.*valuePt1xPt2)/(mMeanPt2))+
				       ((valuePz1SqrSumPz2Sqr+2.*valuePz1xPz2)/(2.*mMeanPz2))+ 
				       (((E1+E2-2.*mMeanE)*(E1+E2-2.*mMeanE))/(2.*(mMeanE2-mMeanE*mMeanE))) ) ));

    double valueFullDen = exp(-1. * (1./double(mNumberOfParticles-1.) * 
				     ( ((valuePt1SqrSumPt2Sqr)/(mMeanPt2)) + 
				       ((valuePz1SqrSumPz2Sqr)/(2.*mMeanPz2)) +
				       (((E1-mMeanE)*(E1-mMeanE) + 
					 (E2-mMeanE)*(E2-mMeanE))/(2.*(mMeanE2 - 
								       mMeanE*mMeanE))) ) ));

  
    mNumeratorFull[mIndex].Fill(Qinv, valueFullNum);
    mDenominatorFull[mIndex].Fill(Qinv, valueFullDen);
  }
}

//_________________
void QinvCorrFctnKt_MomAndEnCons::AddMixedPair(const StHbtPair* pair) {

  int mIndex = (int)(((fabs(pair->kT())-mKtMin)/mDeltaKt));
  if ((mIndex>=0)&&(mIndex<mNumberCFs)) {
    double Qinv = fabs(pair->qInv());
    mDenominator->Fill(Qinv, 1.);
  }
}
