#include "HbtFitter.h"
#include <TMath.h>

//_________________
HbtFitter::HbtFitter(Int_t PartType, Int_t SourceForm,
                     Int_t NonFemtoForm, Bool_t ExcludePoints) {

  //Set basic parameters
  SetPartType(PartType);
  SetSourceForm(SourceForm);
  SetNonFemtoForm(NonFemtoForm);
  SetExcludeFromFit(ExcludePoints);
  mBckgFitIsDone = false;

  mQinvMaxVal = 0.8;
  mQoutMaxVal = 0.8;
  mQsideMaxVal = 0.8;
  mQlongMaxVal = 0.8;

  mQinvNbins = 40;
  mQoutNbins = 40;
  mQsideNbins = 40;
  mQlongNbins = 40;

  //Initialize 1D non-femto parameters
  for(Int_t iIter=0; iIter<3; iIter++) {
    mBckg1DParArr[iIter] = 0;
    mBckg1DParErrArr[iIter] = 0;
  }

  //Initialize 3D non-femto parameters
  for(Int_t iIter1=0; iIter1<3; iIter1++) {  //Out,Side,Long
    //iIter1 corresponds to out,side,long
    for(Int_t iIter2=0; iIter2<3; iIter2++) {
      //iIter2 corresponds to the mBckgParNum
      mBckg3DParArr[iIter1][iIter2] = 0;
      mBckg3DParErrArr[iIter1][iIter2] = 0;
    }
  }
}

//_________________
HbtFitter::~HbtFitter() {
  /* empty */
}

//_________________
void HbtFitter::SetSourceForm(Int_t SourceForm) {

  /*
  if(SourceForm<0 || mSourceForm>1) {
  cout << "Unknown source shape: " << SourceForm << ". Resetting to 0" << endl;
    SourceForm = 0;
  }
  */
  mSourceForm = SourceForm;
  //cout << "Source shape: " << mSourceForm << endl;
}

//_________________
void HbtFitter::SetNonFemtoForm(Int_t NonFemtoForm) {

  /*
    Non-femtoscopic forms:
    0 - constant
    1 - linear
    2 - 2nd oreder polynomial
    3 - sqrt(1 + a x^2 + b x^4)  (sqrt tail)
    4 - a ( 1 + e^{b x^2})       (Gaussian tail)
    5 - a / (1 + b x^2)          (hyperbolic tail)
   */

  mNonFemtoForm = NonFemtoForm;

  switch(mNonFemtoForm) {
  case 0 :
    mBckgParNum = 1;
    break;
  case 1 :
    mBckgParNum = 2;
    break;
  case 2 :
    mBckgParNum = 3;
    break;
  case 3 :
    mBckgParNum = 3;
    break;
  case 4 :
    mBckgParNum = 2;
    break;
  case 5 :
    mBckgParNum = 2;
    break;
  default :
#ifdef MESSAGING
    std::cout << "HbtFitter::SetNonFemtoForm  Unknown non-femtoscopic form: "
	      << mNonFemtoForm << " Switching to constant." << std::endl;
#endif
    mBckgParNum = 1;
    break;
  }

}

//_________________
Double_t HbtFitter::FitFunction(double *x, double *par) {

  Double_t mRetVal = 0;

  //Double check if the points were really exluded
  if(mExcludeRange == 1) {
    if(x[0] > mExcludedVal[0] && x[0]<mExcludedVal[1]) {
      TF1::RejectPoint();
      return 0;
    }
  }

  switch (mNonFemtoForm) {
  case 0 : 
    mRetVal = BowlerSinyukovProcedure(x,par) * ConstantTail(x, &par[2]);
    break;
  case 1 :
    mRetVal = BowlerSinyukovProcedure(x,par) * LinearTail(x,&par[2]);
    break;
  case 2 :
    mRetVal = BowlerSinyukovProcedure(x,par) * Pol2Tail(x,&par[2]);
    break;
  case 3 :
    mRetVal = BowlerSinyukovProcedure(x,par) * SqrtTail(x,&par[2]);
    break;
  case 4 :
    mRetVal = BowlerSinyukovProcedure(x,par) * GaussianTail(x,&par[2]);
    break;
  case 5 :
    mRetVal = BowlerSinyukovProcedure(x,par) * HyperbolicTail(x,&par[2]);
    break;
  default : 
    mRetVal = BowlerSinyukovProcedure(x,par) * ConstantTail(x, &par[2]);
#ifdef MESSAGING
    std::cout << "Unknown form of non-femtoscopic correlations: "  
	      << mNonFemtoForm << std::endl;
#endif
    break;
  }

  return mRetVal;
}

//_________________
Double_t HbtFitter::FitFunction3D(double *x, double *par) {

  Double_t mRetVal = 0;
  
  switch(mNonFemtoForm) {
  case 0 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * ConstantTail3D(x,&par[4]);
    break;
  case 1 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * LinearTail3D(x,&par[4]);
    break;
  case 2 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * Pol2Tail3D(x,&par[4]);
    break;
  case 3 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * SqrtTail3D(x,&par[4]);
    break;
  case 4 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * GaussianTail3D(x,&par[4]);
    break;
  case 5 :
    mRetVal = BowlerSinyukovProcedure3D(x,par) * HyperbolicTail(x,&par[4]);
    break;
  default :
#ifdef MESSAGING    
    std::cout << "FitFunction::Unknown form of non-femtoscopic correlations: "  
	      << mNonFemtoForm << " Switching to constant." << std::endl;
#endif
    mRetVal = BowlerSinyukovProcedure3D(x,par) * ConstantTail(x,&par[4]);
    break;
  }

  return mRetVal;
}

//_________________
Double_t HbtFitter::BowlerSinyukovProcedure(double *x, double *par) {

  Double_t mCoulombVal = 0;
  Double_t mStepQinv = mQinvMaxVal/mQinvNbins;
  Double_t mCorrVal = 0;

  Int_t mQinvBin = (Int_t)( (*x)/mStepQinv );

  //Retrieving Coulomb correction for the particle hypothesis
  switch(mPartType) {
  case 3 :
    mCoulombVal = PI_LONG[mQinvBin];
    break;
  case 4 : 
    mCoulombVal = KK_LONG[mQinvBin];
    break;
  default :
#ifdef MESSAGING
    std::cout << "HBTFitter::BowlerSinyukovProcedure(x,par)" 
	      << " Impossible particle type: " << mPartType << std::endl;
    std::cout << "Coulomb correction is ignored" << std::endl;
#endif
    mCoulombVal = 1;
    break;
  }

  switch(mSourceForm) {
  case 0 : //Gassian
    mCorrVal = (1 - par[0] + par[0] * mCoulombVal * (1 + 
	        TMath::Exp( (-1.)*par[1]*par[1]*(*x)*(*x) ) ) );
    break;
  case 1 : //Exponential
    mCorrVal = (1 - par[0] + par[0] * mCoulombVal * (1 + 
		TMath::Exp( (-1.)*par[1]*(*x) ) ) );
    break;
  default : //Gaussian without Coulomb correction
    mCorrVal = (1 + par[0] * TMath::Exp((-1.)*par[1]*par[1]*(*x)*(*x) ) );
#ifdef MESSAGING
    std::cout << "Unknown source form. Switching to the Goldhaber parameterization." 
	      << std::endl;
#endif
  }

  return mCorrVal;
}

//_________________
Double_t HbtFitter::BowlerSinyukovProcedure3D(double *x, double *par) {

  Double_t mCoulombVal = 1.;
  Double_t mStepQinv = mQinvMaxVal/mQinvNbins;   // 1./50 = 0.02 GeV/c

  Double_t mStepOut = mQoutMaxVal/mQoutNbins;
  Double_t mStepSide = mQsideMaxVal/mQsideNbins;
  Double_t mStepLong = mQlongMaxVal/mQlongNbins;


  Int_t mOutBin, mSideBin, mLongBin;
  mOutBin  = (Int_t)( x[0]/mStepOut );
  mSideBin = (Int_t)( x[1]/mStepSide );
  mLongBin = (Int_t)( x[2]/mStepLong );
#ifdef MESSAGING
  std::cout << "HBTFitter::BowlerSinyukovProcedure(x,y,z,par)"
	    << " Fit the bins... QoutBin: " << mOutBin 
	    << " QsideBin: " << mSideBin << " QlongBin: " << mLongBin
	    << std::endl;
#endif

  Double_t mCorrVal = 0;
  Double_t mQinvWeight=0;
  Int_t mCoulombBin=49;

  //Retrieving the Qinv weight from the 3D histo
  if(mPartType==3 || mPartType==4) {
    Double_t mQinvWeight = hBP3DCoulWeightQinv->GetBinContent(mOutBin,mSideBin,mLongBin);
    mCoulombBin = (Int_t)(mQinvWeight/mStepQinv);
#ifdef MESSAGING
    std::cout << "QinvWeight: " << mQinvWeight 
	      << " mCoulombBin: " << mCoulombBin << std::endl;
#endif
  }

  switch(mPartType) {
  case 3 :
    mCoulombVal = PI_LONG[mCoulombBin];
    break;
  case 4 : 
    mCoulombVal = KK_LONG[mCoulombBin];
#ifdef MESSAGING    
    std::cout << "Kaon Coulomb correction val: " << mCoulombVal << std::endl;
#endif
    break;
  default :
#ifdef MESSAGING
    std::cout << "HBTFitter::BowlerSinyukovProcedure(x,y,z,par)" 
	      << " Impossible particle type: " << mPartType << std::endl;
    std::cout << "Coulomb correction is ignored" << std::endl;
#endif
    mCoulombVal = 1;
    break;
  }

  switch (mSourceForm) {
  case 0 : //Gaussian source
    mCorrVal = (1 - par[0] + par[0] * mCoulombVal *
		(1 + TMath::Exp( (-1.)*par[1]*par[1]*x[0]*x[0] +
				 (-1.)*par[2]*par[2]*x[1]*x[1] + 
				 (-1.)*par[3]*par[3]*x[2]*x[2] ) ) );
    break;
  case 1 : //Exponential source
    mCorrVal = (1 - par[0] + par[0] * mCoulombVal *
		(1 + TMath::Exp( (-1.)*par[1]*x[0] +
				 (-1.)*par[2]*x[1] +
				 (-1.)*par[3]*x[2] ) ) );
    break;
  default :
#ifdef MESSAGING
    std::cout << "Improssible source type: " << mSourceForm 
	      << " Switching to the Goldhaber parameterization." 
	      << std::endl;
#endif

    mCorrVal = (1 + par[0] * TMath::Exp( (-1.)*par[1]*par[1]*x[0]*x[0] +
					 (-1.)*par[2]*par[2]*x[1]*x[1] + 
					 (-1.)*par[3]*par[3]*x[2]*x[2] ) ) ;
    break;
  }

  return mCorrVal;
}

//_________________
Double_t HbtFitter::ConstantTail(double *x, double *par) {

  Double_t mRetVal = par[0];
  return mRetVal;
}

//_________________
Double_t HbtFitter::ConstantTail3D(double *x, double *par) {

  Double_t mRetVal = par[0];
  return mRetVal;
}

//_________________
Double_t HbtFitter::LinearTail(double *x, double *par) {

  Double_t mRetVal = par[0]*(1 + par[1]*(*x));
  return mRetVal;
}

//_________________
Double_t HbtFitter::LinearTail3D(double *x, double *par) {

  Double_t mRetVal = par[0]*(1 + par[1]*x[0] + par[2]*x[1] + par[3]*x[2]);
  return mRetVal;
}

//_________________
Double_t HbtFitter::Pol2Tail(double *x, double *par) {

  Double_t mRetVal = par[0]*(1 + par[1]*(*x) + par[2]*(*x)*(*x));
  return mRetVal;
}

//_________________
Double_t HbtFitter::Pol2Tail3D(double *x, double *par) {

  Double_t mRetVal = par[0]*(1 + par[1]*x[0] + par[2]*x[1] + par[3]*x[2] +
			     par[4]*x[0]*x[0] + par[5]*x[1]*x[1] + par[6]*x[2]*x[2]);
  return mRetVal;
}

//_________________
Double_t HbtFitter::GaussianTail(double *x, double *par) {

  Double_t mRetVal = par[0] * (1 + TMath::Exp(-par[1]*(*x)*(*x)));
  return mRetVal;
}

//_________________
Double_t HbtFitter::GaussianTail3D(double *x, double *par) {

  Double_t mRetVal = par[0] * ( 1 + TMath::Exp( -par[1]*x[0]*x[0] ) + 
				TMath::Exp( -par[2]*x[1]*x[1] ) +
				TMath::Exp( -par[3]*x[2]*x[2] ) );
  return mRetVal;
}

//_________________
Double_t HbtFitter::SqrtTail(double *x, double *par) {

  Double_t mRetVal = par[0] * TMath::Sqrt(1 + par[1]*(*x)*(*x) + par[2]*(*x)*(*x)*(*x)*(*x));
  return mRetVal;
}

//_________________
Double_t HbtFitter::SqrtTail3D(double *x, double *par) {

  Double_t mRetVal = par[0] * TMath::Sqrt( 1 + par[1]*x[0]*x[0] + par[2]*x[1]*x[1] + 
					   par[3]*x[2]*x[2] +
					   par[4]*x[0]*x[0]*x[0]*x[0] + par[5]*x[1]*x[1]*x[1]*x[1] +
					   par[6]*x[2]*x[2]*x[2]*x[2] );
  return mRetVal;
}

//________________
Double_t HbtFitter::HyperbolicTail(double *x, double *par) {
  Double_t mRetVal = par[0] / (1 + par[1]*(*x)*(*x));
  return mRetVal;
}

//________________
Double_t HbtFitter::HyperbolicTail3D(double *x, double *par) {

  Double_t mRetVal = par[0] / (1 + par[1]*x[0]*x[0] + par[2]*x[1]*x[1] + par[3]*x[2]*x[2] );
  return mRetVal;
}

//_________________
Double_t HbtFitter::FitBackgroundFunction(double *x, double *par) {

  //Fit the baseline distribution only
  Double_t mRetVal = 0;
  switch(mNonFemtoForm) {
  case 0 :
    mRetVal = ConstantTail(x,par);
    break;
  case 1 :
    mRetVal = LinearTail(x, par);
    break;
  case 2 :
    mRetVal = Pol2Tail(x, par);
    break;
  case 3 :
    mRetVal = SqrtTail(x, par);
    break;
  case 4 :
    mRetVal = GaussianTail(x, par);
    break;
  case 5 :
    mRetVal = HyperbolicTail(x, par);
    break;
  default :
    mRetVal = ConstantTail(x,par);
    break;
  }

  mBckgFitIsDone = true;
  return mRetVal;
}

//_________________
Double_t HbtFitter::FitBackgroundFunction3D(double *x, double *par) {

  Double_t mRetVal = 0;
  switch(mNonFemtoForm) {
  case 0 :
    mRetVal = ConstantTail3D(x,par);
    break;
  case 1 :
    mRetVal = LinearTail3D(x,par);
    break;
  case 2 :
    mRetVal = Pol2Tail3D(x,par);
    break;
  case 3 :
    mRetVal = SqrtTail3D(x,par);
    break;
  case 4 :
    mRetVal = GaussianTail3D(x,par);
    break;
  case 5 :
    mRetVal = HyperbolicTail3D(x,par);
    break;
  default :
    mRetVal = ConstantTail3D(x,par);
    break;
  }

  mBckgFitIsDone = true;
  return mRetVal;
}



//_________________
void HbtFitter::Set3DBinsNum(Int_t OutNBins, Int_t SideNBins, Int_t LongNBins) {
  mQoutNbins = OutNBins;
  mQsideNbins = SideNBins;
  mQlongNbins = LongNBins;
}

//_________________
void HbtFitter::Set3DMaxVal(Double_t MaxOut, Double_t MaxSide, Double_t MaxLong) {
  mQoutMaxVal = MaxOut;
  mQsideMaxVal = MaxSide;
  mQlongMaxVal = MaxLong;
}

//_________________
void HbtFitter::Set1DBackgroundParameters(vector<double> par) {
  
  for(UInt_t i=0; i<par.size(); i++) {
    mBckg1DParArr[i] = par.at(i);
  }
}

//_________________
void HbtFitter::Set3DBackgroundParameters(vector<double> parOut,
					  vector<double> parSide,
					  vector<double> parLong) {
  //Assuming the same amount of parameters 
  //for out, side and long components
  for(UInt_t i=0; i<parOut.size(); i++) {
    mBckg3DParArr[0][i] = parOut.at(i);
    mBckg3DParArr[1][i] = parSide.at(i);
    mBckg3DParArr[2][i] = parLong.at(i);
  }
}

//_________________
void HbtFitter::SetExcludeRange(Double_t xMin, Double_t xMax) {

  //Exlude range from the fit
  mExcludedVal[0] = xMin;
  mExcludedVal[1] = xMax;
}

//_________________
Double_t HbtFitter::FitSignalFunction(double *x, double *par) {

  Double_t mCorrVal = 0;

  switch(mNonFemtoForm) {
  case 0 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2];
    break;
  case 1 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2] * (1 + mBckg1DParArr[1]*(*x));
    break;
  case 2 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2] * (1 + mBckg1DParArr[1]*(*x) + 
							   mBckg1DParArr[2]*(*x)*(*x));
    break;
  case 3 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2] * TMath::Sqrt(1 + 
               mBckg1DParArr[1]*(*x)*(*x) + mBckg1DParArr[2]*(*x)*(*x)*(*x)*(*x));
    break;
  case 4 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2] * (1 + 
	       TMath::Exp( (-1.)*mBckg1DParArr[1]*(*x)*(*x) ) );
    break;
  case 5 :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2] / (1 + mBckg1DParArr[1]*(*x)*(*x) );
    break;
  default :
    mCorrVal = BowlerSinyukovProcedure(x,par) * par[2];
    //break;
  }

  return mCorrVal;
}

//_________________
Double_t HbtFitter::FitSignalFunction3D(double *x, double *par) {

  Double_t mCorrVal = 0;

  switch(mNonFemtoForm) {
  case 0 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4];
    break;
  case 1 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4] * (1 + mBckg3DParArr[0][1]*x[0] +
							       mBckg3DParArr[1][1]*x[1] +
							       mBckg3DParArr[2][1]*x[2]);
    break;
  case 2 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4] * 
      (1 + mBckg3DParArr[0][1]*x[0] + mBckg3DParArr[0][2]*x[0]*x[0] +
       mBckg3DParArr[1][1]*x[1] + mBckg3DParArr[1][2]*x[1]*x[1] +
       mBckg3DParArr[2][1]*x[2] + mBckg3DParArr[2][2]*x[2]*x[2]);
    break;
  case 3 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4] * 
      TMath::Sqrt(1 + mBckg3DParArr[0][1]*x[0]*x[0] + mBckg3DParArr[0][2]*x[0]*x[0]*x[0]*x[0] +
		  mBckg3DParArr[1][1]*x[1]*x[1] + mBckg3DParArr[1][2]*x[1]*x[1]*x[1]*x[1] +
		  mBckg3DParArr[2][1]*x[2]*x[2] + mBckg3DParArr[2][2]*x[2]*x[2]*x[2]*x[2]);
    break;
  case 4 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4] * 
      (1 + TMath::Exp( (-1.)*mBckg3DParArr[0][1]*x[0]*x[0] ) +
       TMath::Exp( (-1.)*mBckg3DParArr[1][1]*x[1]*x[1] ) +
       TMath::Exp( (-1.)*mBckg3DParArr[2][1]*x[2]*x[2] ) );
    break;
  case 5 :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4] / (1 + mBckg3DParArr[0][1]*x[0]*x[0] +
							       mBckg3DParArr[1][1]*x[1]*x[1] + 
							       mBckg3DParArr[2][1]*x[2]*x[2]);
    break;
  default :
    mCorrVal = BowlerSinyukovProcedure3D(x,par) * par[4];
    //break;
  }

  return mCorrVal;
}


