#include "StHbtMaker/CorrFctn/QinvCorrFctnSmeared.h"
#include <TMath.h>
#include <TRandom3.h>
#include <TString.h>
#include <cstdio>
#include <sstream>

//Richard's FSI weights
#include "StHbtMaker/ThLednickyWeight/WLedCOMMONS.h"
extern "C" void fsiini_();
extern "C" void ltran12_();
extern "C" void fsiw_();
extern LEDWEIGHTCommon LEDWEIGHT;
extern MOMLABCommon MOMLAB;
extern FSIMOMCommon FSIMOM;
extern FSINSCommon FSINS;
extern FSICOORCommon FSICOOR;
extern FSIPOCCommon FSIPOC;
extern FSIPRFCommon FSIPRF;

#ifdef __ROOT__
ClassImp(QinvCorrFctnSmeared)
#endif

//_________________
QinvCorrFctnSmeared::QinvCorrFctnSmeared(char* title,
					 const int& nbins, 
					 const float& QinvLo, 
					 const float& QinvHi){

  //Set up qInv functions
  mMixedQinv = new StHbt1DHisto;
  mMixedQinvWeight = new StHbt1DHisto;
  mMixedQinvSmeared = new StHbt1DHisto;
  mMixedQinvSmearedWeight = new StHbt1DHisto;

  stringstream sMQ, sMQW, sMQS, sMQSW, sTitle;
  sMQ.str(""); sMQW.str(""); sMQS.str(""); sMQSW.str(""); sTitle.str("");
  sMQ << "MQ" << title;
  sMQW << "MQW" << title;
  sMQS << "MQS" << title ;
  sMQSW << "MQSW" << title;

  sTitle.str("");
  sTitle << sMQ.str();
  mMixedQinv->SetName(sTitle.str().c_str());
  mMixedQinv->SetTitle(sTitle.str().c_str());
  mMixedQinv->SetBins(nbins,QinvLo,QinvHi);
  mMixedQinv->SetDirectory(0);
  mMixedQinv->Sumw2();
  
  sTitle.str("");
  sTitle << sMQW.str();
  mMixedQinvWeight->SetName(sTitle.str().c_str());
  mMixedQinvWeight->SetTitle(sTitle.str().c_str());
  mMixedQinvWeight->SetBins(nbins,QinvLo,QinvHi);
  mMixedQinvWeight->SetDirectory(0);
  mMixedQinvWeight->Sumw2();

  sTitle.str("");
  sTitle << sMQS.str();
  mMixedQinvSmeared->SetName(sTitle.str().c_str());
  mMixedQinvSmeared->SetTitle(sTitle.str().c_str());
  mMixedQinvSmeared->SetBins(nbins,QinvLo,QinvHi);
  mMixedQinvSmeared->SetDirectory(0);
  mMixedQinvSmeared->Sumw2();

  sTitle.str("");
  sTitle << sMQSW.str();
  mMixedQinvSmearedWeight->SetName(sTitle.str().c_str());
  mMixedQinvSmearedWeight->SetTitle(sTitle.str().c_str());
  mMixedQinvSmearedWeight->SetBins(nbins,QinvLo,QinvHi);
  mMixedQinvSmearedWeight->SetDirectory(0);
  mMixedQinvSmearedWeight->Sumw2();
  
  //Varibales that can be modified via methods
  //The default values are choosen for AuAu2011
  mPt[0] = 0.0023;
  mPt[1] = 0.0044;
  mPt[2] = 0.00362;
  mPtAlpha = -0.73;

  mTheta[0] = 0.000676;
  mTheta[1] = 0.000604;
  mThetaAlpha = -1.133;

  mPhi[0] = 0.000939;
  mPhi[1] = 0.000503;
  mPhiAlpha = -1.389;
  
  mA[0] = 0.00129;
  mA[1] = -0.00059;
  mA[2] = -1.53;

  mB[0] = 0.000003;
  mB[1] = -0.000002;
  mB[2] = -2.1;

  mC[0] = 0.000025;
  mC[1] = -0.000008;
  mC[2] = -2.1;

  //FSI parameters
  mITest = 1;
  mCalcCoul = 1;
  mCalcQS = 1;
  mCalcStrong = 0;
  mPartCoulomb = 15; //K+K+

  mMass1 = 0.493677;
  mMass2 = 0.493677;
  mCharge1 = 1;
  mCharge2 = 1;

  //Radii and lambda
  mLambda = 0.6;
  mRadii[0] = 3.;
  mRadii[1] = 3.;
  mRadii[2] = 3.;
}

//_________________
QinvCorrFctnSmeared::~QinvCorrFctnSmeared() {
  delete mMixedQinv;
  delete mMixedQinvWeight;
  delete mMixedQinvSmeared;
  delete mMixedQinvSmearedWeight;
}

//_________________
void QinvCorrFctnSmeared::Finish() {
  /* nothing to do */
}

//_________________
void QinvCorrFctnSmeared::AddRealPair(const StHbtPair* pair) {
  /* nothing to do */
}

//_________________
StHbtString QinvCorrFctnSmeared::Report() {
  int mMQentries=0;
  int mMQWentries=0;
  int mMQSentries=0;
  int mMQSWentries=0;
  
  mMQentries = (int)mMixedQinv->GetEntries();
  mMQWentries = (int)mMixedQinvWeight->GetEntries();
  mMQSentries = (int)mMixedQinvSmeared->GetEntries();
  mMQSWentries = (int)mMixedQinvSmearedWeight->GetEntries();

  string stemp = "Qinv Correlation Function Smeared report:\n";
  char ctemp[100];
  sprintf(ctemp,"Number of entries in all mixed events:\t%i\n",mMQentries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in all mixed weighted events:\t%i\n",
	  mMQWentries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in all mixed smeared events:\t%i\n",
	  mMQSentries);
  stemp += ctemp;
  sprintf(ctemp,"Number of entries in all mixed smeared weighted events:\t%i\n",
	  mMQSWentries);
  stemp += ctemp;
  StHbtString returnThis = stemp;
  return returnThis;
}

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

  //FSI initialization
  LEDWEIGHT.ITEST = mITest;
  FSINS.NS = 4;
  FSINS.ICH = mCalcCoul;
  FSINS.ISI = mCalcStrong;
  FSINS.IQS = mCalcQS;
  FSINS.I3C = 0;
  FSINS.LL = mPartCoulomb;

  FSIPOC.AM1 = mMass1;
  FSIPOC.C1 = mCharge1;
  FSIPOC.AM2 = mMass2;
  FSIPOC.C2 = mCharge2;

  fsiini_();

  double px1 = pair->track1()->FourMomentum().x();
  double py1 = pair->track1()->FourMomentum().y();
  double pz1 = pair->track1()->FourMomentum().z();
  double e1 = pair->track1()->FourMomentum().e();
  double theta1 = pair->track1()->FourMomentum().theta();
  double massSqr1 = e1*e1 - px1*px1 - py1*py1 - pz1*pz1;
  double pT1 = sqrt(px1*px1 + py1*py1);

  double px2 = pair->track2()->FourMomentum().x();
  double py2 = pair->track2()->FourMomentum().y();
  double pz2 = pair->track2()->FourMomentum().z();
  double e2 = pair->track2()->FourMomentum().e();
  double theta2 = pair->track2()->FourMomentum().theta();
  double massSqr2 = e2*e2 - px2*px2 - py2*py2 - pz2*pz2;
  double pT2 = sqrt(px2*px2 + py2*py2);

  double dRout1 = gRandom->Gaus(0.,mRadii[0]);
  double dRside1 = gRandom->Gaus(0., mRadii[1]);
  double dRlong1 = gRandom->Gaus(0., mRadii[2]);
  double dT1 = 0.;

  double dRout2 = gRandom->Gaus(0.,mRadii[0]);
  double dRside2 = gRandom->Gaus(0., mRadii[1]);
  double dRlong2 = gRandom->Gaus(0., mRadii[2]);
  double dT2 = 0.;

  FSIMOM.P1X = px1;
  FSIMOM.P1Y = py1;
  FSIMOM.P1Z = pz1;
  FSIMOM.E1 = e1;

  FSIMOM.P2X = px2;
  FSIMOM.P2Y = py2;
  FSIMOM.P2Z = pz2;
  FSIMOM.E2 = e2;

  //Boost to the lab frame
  double Px = px1 + px2;
  double Py = py1 + py2;
  double Pz = pz1 + pz2;
  double E = e1 + e2;

  double Pt = Px*Px + Py*Py;
  double Mt = sqrt(E*E - Pz*Pz);    //Should be sqrt(E^2-Pz^2)
  double Minv = sqrt(E*E - Px*Px - Py*Py - Pz*Pz);
  double BetaT = Pt/Mt;
  double GammaT = Mt/Minv;

  double Rout1_LCMS = dRout1;
  double Rout2_LCMS = dRout2;
  double T1_LCMS = GammaT*(dT1 + BetaT*dRout1);
  double T2_LCMS = GammaT*(dT2 + BetaT*dRout2);
  double BetaZ = Pz/E;
  double GammaZ = E/Mt;

  double Rout1_lab = Rout1_LCMS;
  double Rout2_lab = Rout2_LCMS;
  double Rlong1_lab = GammaZ*(dRlong1 + BetaZ*T1_LCMS);
  double Rlong2_lab = GammaZ*(dRlong2 + BetaZ*T2_LCMS);
  double T1_lab = GammaZ*(T1_LCMS + BetaZ*dRlong1);
  double T2_lab = GammaZ*(T2_LCMS + BetaZ*dRlong2);

  double t1 = T1_lab;
  double x1 = Rout1_lab;
  double y1 = dRside1;
  double z1 = Rlong1_lab;

  double t2 = T2_lab;
  double x2 = Rout2_lab;
  double y2 = dRside2;
  double z2 = Rlong2_lab;

  //Set emission points
  FSICOOR.X1 = x1;
  FSICOOR.Y1 = y1;
  FSICOOR.Z1 = z1;
  FSICOOR.T1 = t1;

  FSICOOR.X2 = x2;
  FSICOOR.Y2 = y2;
  FSICOOR.Z2 = z2;
  FSICOOR.T2 = t2;

  //Perform boost to PRF
  ltran12_();
 
  //Get CF's weight according to Richard Lednicky
  fsiw_();
  double LedWeight = float(LEDWEIGHT.WEIN);
  
  //Smear pT, theta and phis
  double dpT1 = gRandom->Gaus( mA[0]+mA[1]*pow(pT1,mA[2]),
			       mPt[0]+mPt[1]*pow(pT1,mPtAlpha)+mPt[2]*pT1 );
  double dpT2 = gRandom->Gaus( mA[0]+mA[1]*pow(pT2,mA[2]),
			       mPt[0]+mPt[1]*pow(pT2,mPtAlpha)+mPt[2]*pT2 );

  double dtheta1 = gRandom->Gaus( mB[0]+mB[1]*pow(pT1,mB[2]),
				  mTheta[0]+mTheta[1]*pow(pT1,mThetaAlpha) );
  double dtheta2 = gRandom->Gaus( mB[0]+mB[1]*pow(pT2,mB[2]),
				  mTheta[0]+mTheta[1]*pow(pT2,mThetaAlpha) );

  double dphi1 = gRandom->Gaus( mC[0]+mC[1]*pow(pT1,mC[2]),
				mPhi[0]+mPhi[1]*pow(pT1,mPhiAlpha) );
  double dphi2 = gRandom->Gaus( mC[0]+mC[1]*pow(pT2,mC[2]),
				mPhi[0]+mPhi[1]*pow(pT2,mPhiAlpha) );

  //Smeared four-momentum components
  double px1s = px1*(1. + dpT1) - py1*dphi1;
  double py1s = py1*(1. + dpT1) - px1*dphi1;
  double pz1s = pz1*(1. + dpT1) - pT1*dtheta1/(sin(theta1)*sin(theta1));
  double pe1s = sqrt(massSqr1 + px1s*px1s + py1s*py1s + pz1s*pz1s);

  double px2s = px2*(1. + dpT2) - py2*dphi2;
  double py2s = py2*(1. + dpT2) - px2*dphi2;
  double pz2s = pz2*(1. + dpT2) - pT2*dtheta2/(sin(theta2)*sin(theta2));
  double pe2s = sqrt(massSqr2 + px2s*px2s + py2s*py2s + pz2s*pz2s);

  //Make new transform to PRF
  FSIMOM.P1X = px1s;
  FSIMOM.P1Y = py1s;
  FSIMOM.P1Z = pz1s;
  FSIMOM.E1 = pe1s;

  FSIMOM.P2X = px2s;
  FSIMOM.P2Y = py2s;
  FSIMOM.P2Z = pz2s;
  FSIMOM.E2 = pe2s;
  ltran12_();
  
  double ak = 2 * FSIPRF.AK;

  //Get actual CF weigth
  double weight = 1. - mLambda + mLambda*LedWeight;
  
  //Fill histograms
  mMixedQinv->Fill(fabs(pair->qInv()), 1.);
  mMixedQinvWeight->Fill(fabs(pair->qInv()), weight);
  mMixedQinvSmeared->Fill(ak, 1.);
  mMixedQinvSmearedWeight->Fill(ak, weight);
  
}
