#include "StFemtoDstPythia6Maker.h"
#include "TBranch.h"
#include "TMath.h"
#include "TLorentzVector.h"
#include "TClonesArray.h"
#include "TObjArray.h"
#include "TDatabasePDG.h"

ClassImp(StFemtoDstPythia6Maker)

#define MAXFILESIZE 1900000000
const Float_t PYTHIA_BFIELD = 4.510600000;  
const Float_t PYTHIA_RBFIELD = -4.510600000;

//_________________
StFemtoDstPythia6Maker::StFemtoDstPythia6Maker(TPythia6 *pythiaDst,
					       const Char_t *oFileName) {

  std::cout << "StFemtoDstPythia6Maker::StFemtoDstPythia6Maker - Creating an instance...";
  mOutFileName = oFileName;
  mPythiaDst = pythiaDst;
  mCompression = 9;

  mEventIsGood = false;
  mNEventsIn = 0;
  mIsGoodTrack = false;

  mIsFullField = false;

  //Initialize single-particle cut variables
  mTrackMomentum[0] = 0.15;
  mTrackMomentum[1] = 1.75;
  mTrackDca[0] = 0.;
  mTrackDca[1] = 3.;
  mTrackDcaGlobal[0] = 0.;
  mTrackDcaGlobal[1] = 3.;
  mTrackEta[0] = -1.1;
  mTrackEta[1] = 1.1;

  //Flags and cuts for exclusion
  mRemovePions = false;
  mRemoveKaons = false;
  mRemoveProtons = true;
  randy = new TRandom3(0);

  std::cout << "\t[DONE]" << std::endl;
}

//_________________
StFemtoDstPythia6Maker::~StFemtoDstPythia6Maker() {
  /* nothing to do */
}

//_________________
Int_t StFemtoDstPythia6Maker::Init() {

  std::cout << "StFemtoDstPythia6Maker::Init - Initializing the maker"
	    << std::endl 
	    << "Creating output file: " << mOutFileName;

  mFemtoEvent = new StFemtoEvent();
  mOutFile = new TFile(mOutFileName, "RECREATE");
  mOutFile->SetCompressionLevel(mCompression);
  std::cout << "\t[DONE]" << std::endl;

  mTree = new TTree("StFemtoDst","StFemtoDst");
  mTree->SetMaxTreeSize(MAXFILESIZE);
  mTree->Branch("StFemtoEvent","StFemtoEvent", &mFemtoEvent);

  mNBytes = 0;
  std::cout << "StFemtoDstPythia6Maker::Init - Initialization has been finished"
	    << std::endl;

  return StMaker::Init();
}

//_________________
void StFemtoDstPythia6Maker::Clear(Option_t *option) {
  StMaker::Clear();
}

//_________________
Int_t StFemtoDstPythia6Maker::Make() {

  Float_t mMagneticField;
  if(mIsFullField) {
    mMagneticField = PYTHIA_BFIELD;
  }
  else {
    mMagneticField = PYTHIA_RBFIELD;
  }

  TObjArray *pyParticles = NULL;
  TMCParticle *pyParticle = NULL;

  Int_t mNumParticles = 0;
  Int_t pyPosTrkNum = 0;
  Int_t pyPosTrk2Num = 0;
  //Int_t pyNegTrkNum = 0;
  Int_t pyRefMult = 0;
  Int_t pyRefMult2 = 0;
  Int_t pyPrimTracksNum = 0;
  Int_t pyGlobTracksNum = 0;
  
  Int_t pyKF, pyKS, pyCharge;
  Float_t pyTrkMass;
  Float_t pyTrkNSigmaElectron, pyTrkNSigmaPion;
  Float_t pyTrkNSigmaKaon, pyTrkNSigmaProton;

  Float_t pyPx, pyPy, pyPz, pyP, pyEta;
  Float_t pyVx, pyVy, pyVz, pyDcaXY;

  //Generate non-empty event
  do {
    CleanVariables();
    mPythiaDst->GenerateEvent();
    pyParticles = mPythiaDst->GetListOfParticles();
    if(pyParticles) {
      mEventIsGood = true;
    }
  } while(!mEventIsGood);

  mNEventsIn++;

  mNumParticles = pyParticles->GetEntries();

  //mPythiaDst->Pylist(1);

  //Number of particles in acceptance
  pyRefMult = 0;

  StFemtoTrack *mFTrack = NULL;

  //Simulated track loop
  for(Int_t iTrk=0; iTrk<mNumParticles; iTrk++) {

    pyParticle = (TMCParticle*)pyParticles->At(iTrk);
    mIsGoodTrack = false;
    pyKF = 0;
    pyKS = 0;
    pyKS = pyParticle->GetKS();
    pyCharge = 0;

    //Use final-state particles
    if(!pyParticle || pyKS!=1) continue;

    pyKF = pyParticle->GetKF();
    
    pyCharge = TDatabasePDG::Instance()->GetParticle(pyKF)->Charge();

    pyPx = pyParticle->GetPx();
    pyPy = pyParticle->GetPy();
    pyPz = pyParticle->GetPz();
    pyP = TMath::Sqrt(pyPx*pyPx + pyPy*pyPy + pyPz*pyPz);
    pyEta = 0.5 * TMath::Log( (pyP + pyPz)/ (pyP - pyPz) );

    pyVx = pyParticle->GetVx() * 0.1;
    pyVy = pyParticle->GetVy() * 0.1;
    pyVz = pyParticle->GetVz() * 0.1;
    pyDcaXY = TMath::Sqrt(pyVx*pyVx + pyVy*pyVy);

    //Calculate refMult
    if(TMath::Abs(pyEta)<=0.5 && pyP>=0.075 && pyCharge!=0) {
      pyRefMult++;

      if(pyCharge>0) {
	pyPosTrkNum++;
      }

      /*
      if(pyCharge<0) {
	pyNegTrkNum++;
      }
      */
    }

    //Calculate refMult2
    if(TMath::Abs(pyEta)<=1. && pyP>=0.075 && pyCharge!=0) {
      pyRefMult2++;

      if(pyCharge>0) {
	pyPosTrk2Num++;
      }
    }

    //Calculate number of primary and global tracks
    if(TMath::Abs(pyEta)<=1. && pyP>0.14) {

      pyGlobTracksNum++;

      if(pyDcaXY<=2. && pyVz<=2.) {
	pyPrimTracksNum++;
      }
    }

    //Use pions and kaons only
    if( TMath::Abs(pyKF)!=211 && TMath::Abs(pyKF)!=321 &&
	TMath::Abs(pyKF)!=2212 ) continue;

    if(mRemovePions==true && TMath::Abs(pyKF)==211) continue;
    if(mRemoveKaons==true && TMath::Abs(pyKF)==321) continue;
    if(mRemoveProtons==true && TMath::Abs(pyKF)==2212) continue;

    mIsGoodTrack = AcceptTrack(pyParticle, pyCharge);
    if(!mIsGoodTrack) continue;

    mFTrack = NULL;

    pyTrkMass = pyParticle->GetMass();
    pyCharge = TDatabasePDG::Instance()->GetParticle(pyKF)->Charge();
    if(pyCharge > 0) {
      pyCharge = 1;
    }
    else if(pyCharge < 0) {
      pyCharge = -1;
    }
    Int_t mNHits = 45;
    unsigned int tmap1 = 0;
    unsigned int tmap2 = 0;

    //Set some track parameters
    switch(TMath::Abs(pyKF)) {
    case 211:
      pyTrkNSigmaElectron = -999.;
      pyTrkNSigmaPion = 0.;
      pyTrkNSigmaKaon = -999.;
      pyTrkNSigmaProton = -999.;
      break;
    case 321:
      pyTrkNSigmaElectron = -999.;
      pyTrkNSigmaPion = -999.;
      pyTrkNSigmaKaon = 0.;
      pyTrkNSigmaProton = -999.;
      break;
    case 2212:
      pyTrkNSigmaElectron = -999.;
      pyTrkNSigmaPion = -999.;
      pyTrkNSigmaKaon = -999.;
      pyTrkNSigmaProton = 0.;
      break;
    default:
      pyTrkNSigmaElectron = -999.;
      pyTrkNSigmaPion = -999.;
      pyTrkNSigmaKaon = -999.;
      pyTrkNSigmaProton = -999.;
    }

    Float_t ptot = TMath::Sqrt(pyPx*pyPx + pyPy*pyPy + pyPz*pyPz);
    Float_t energy = TMath::Sqrt(ptot*ptot + pyTrkMass*pyTrkMass);
    mFTrack = mFemtoEvent->AddFemtoTrack();

    mFTrack->SetId( iTrk );
    mFTrack->SetNHits( (Char_t) (pyCharge * mNHits) );
    mFTrack->SetNHitsPoss( mNHits );
    //mFTrack->SetNHitsDedx( mNHits );
    mFTrack->SetNSigmaElectron( pyTrkNSigmaElectron );
    mFTrack->SetNSigmaPion( pyTrkNSigmaPion );
    mFTrack->SetNSigmaKaon( pyTrkNSigmaKaon );
    mFTrack->SetNSigmaProton( pyTrkNSigmaProton );
    mFTrack->SetDedx( dedxMean(pyTrkMass, pyP) );
    //mFTrack->SetChi2( 1. );

    if(ptot < 0.15) { //Tracks only in the inner sector. R=1m
      tmap1 = 0;
      tmap2 = 0;
      for(int iBit=0; iBit<32; iBit++) {
	if(randy->Rndm() < 0.95) {
	  tmap1 |= 1 << iBit;
	}
	else {
	  tmap1 |= 0 << iBit;
	}
      }

      //Default set to the first bit: primary-vertex-used
      tmap1 |= 1 << 0;
    }
    else if(ptot < 0.3) { //Tracks may come to the end of the outer sector: R = 2m
      tmap1 = 0;
      tmap2 = 0;
      for(int iBit=0; iBit<32; iBit++) {

	if(randy->Rndm() < 0.95) {
	  tmap1 |= 1 << iBit;
	}
	else {
	  tmap1 |= 0 << iBit;
	}

	if(iBit < 27) {
	  if(randy->Rndm() < 0.85) {
	    tmap2 |= 1 << iBit;
	  }
	  else {
	    tmap2 |= 0 << iBit;
	  }
	}
      }
      //Default set to the first bit: primary-vertex-used
      tmap1 |= 1 << 0;
    }
    else {
      tmap1 = 0;
      tmap2 = 0;
      for(int iBit=0; iBit<32; iBit++) {

	if(randy->Rndm() < 0.95) {
	  tmap1 |= 1 << iBit;
	}
	else {
	  tmap1 |= 0 << iBit;
	}

	if(iBit < 27) {
	  if(randy->Rndm() < 0.95) {
	    tmap2 |= 1 << iBit;
	  }
	  else {
	    tmap2 |= 0 << iBit;
	  }
	}
      }
      //Default set to the first bit: primary-vertex-used
      tmap1 |= 1 << 0;
    }
    
    mFTrack->SetMap0( tmap1 );
    mFTrack->SetMap1( tmap2 );
    mFTrack->SetP( pyPx, pyPy, pyPz );
    mFTrack->SetDCAxGlobal( pyVx );
    mFTrack->SetDCAyGlobal( pyVy );
    mFTrack->SetDCAzGlobal( pyVz );
    mFTrack->SetGlobalP( pyPx, pyPy, pyPz );
    mFTrack->SetBeta( ptot/energy );
  } //for(Int_t iTrk=0; iTrk<mNumParticles; iTrk++)

  mFemtoEvent->SetEventId( mNEventsIn );
  mFemtoEvent->SetRunId( randy->Uniform(1, 10000000) );
  mFemtoEvent->SetCollisionType( false );
  mFemtoEvent->SetRefMult( pyRefMult );
  mFemtoEvent->SetRefMult2( pyRefMult2 );
  mFemtoEvent->SetRefMultCorr( pyRefMult );
  mFemtoEvent->SetRefMultCorrWeight( 1. );
  mFemtoEvent->SetCent16( 0 );
  mFemtoEvent->SetRefMultPos( pyPosTrkNum );
  mFemtoEvent->SetRefMult2Pos( pyPosTrk2Num );
  //mFemtoEvent->SetRefMultNeg( pyNegTrkNum );
  //mFemtoEvent->SetZDCe( 45000 );
  //mFemtoEvent->SetZDCw( 45000 );
  mFemtoEvent->SetNumberOfBTofHit( pyRefMult );
  mFemtoEvent->SetNumberOfPrimaryTracks( pyPrimTracksNum );
  mFemtoEvent->SetNumberOfGlobalTracks( pyGlobTracksNum );
  mFemtoEvent->SetMagneticField( mMagneticField );
  mFemtoEvent->SetPrimaryVertexPosition(0., 0., 0.);
  mFemtoEvent->SetVpdVz( 0. );
  mFemtoEvent->SetTriggerId( mTriggerId );
  mFemtoEvent->SetPrimaryVertexRanking( 10000 );
  
  if(mEventIsGood) {
    mNBytes += mTree->Fill();
    mFemtoEvent->Clear();
  }

  return kStOk;
}

//_________________
Int_t StFemtoDstPythia6Maker::Finish() {

  mOutFile->cd();
  mOutFile->Write();
  mOutFile->Close();

  std::cout << "****************************************" << std::endl
	    << "StFemtoDstPythia6Maker has been finished" << std::endl
	    << "\t nEventsGenerated: " << mNEventsIn << std::endl
	    << "****************************************" << std::endl;

  return kStOk;
}

//_________________
void StFemtoDstPythia6Maker::CleanVariables() {
  mEventIsGood = false;
}

//_________________
Bool_t StFemtoDstPythia6Maker::AcceptTrack(TMCParticle *trk, Int_t charge) {

  Float_t px = trk->GetPx();
  Float_t py = trk->GetPy();
  Float_t pz = trk->GetPz();
  Float_t p = TMath::Sqrt(px*px + py*py + pz*pz);
  Float_t eta = 0.5*TMath::Log( (p + pz)/ (p - pz) );

  Float_t vx = trk->GetVx() / 10.;
  Float_t vy = trk->GetVy() / 10.;
  Float_t vz = trk->GetVz() / 10.;
  //dcaXY is dcaXYZ now
  Float_t dcaXY = TMath::Sqrt(vx*vx + vy*vy + vz*vz);
  
  Bool_t mIsGoodKinematics = false;
  Bool_t mIsGoodCharge = false;
  Bool_t mIsGoodParticle = false;

  if(charge > 0. || charge < 0.) {
    mIsGoodCharge = true;
  }
 
  if( p >= mTrackMomentum[0] &&
      p <= mTrackMomentum[1] &&
      eta >= mTrackEta[0] &&
      eta <= mTrackEta[1] &&
      dcaXY >= mTrackDca[0] &&
      dcaXY <= mTrackDca[1] &&
      dcaXY >= mTrackDcaGlobal[0] &&
      dcaXY <= mTrackDcaGlobal[1]) {

    mIsGoodKinematics = true;
  }

  if(mIsGoodKinematics == true &&
     mIsGoodCharge == true) {
    mIsGoodParticle = true;
  }
  return mIsGoodParticle;
}

//________________
Double_t StFemtoDstPythia6Maker::dedxMean(Double_t mass, Double_t momentum) {

  Double_t dedxMean;
  Double_t tpcDedxGain = 0.174325e-06;
  Double_t tpcDedxOffset = -2.71889; 
  Double_t tpcDedxRise = 776.626;
  
  Double_t gamma = TMath::Sqrt(momentum*momentum/(mass*mass)+1.);
  Double_t beta = TMath::Sqrt(1. - 1./(gamma*gamma));
  Double_t rise = tpcDedxRise* beta*beta*gamma*gamma;      
  if ( beta > 0)
    dedxMean = tpcDedxGain/(beta*beta) * (0.5*TMath::Log(rise) - beta*beta 
					  - tpcDedxOffset);
  else
    dedxMean = 1000.;
  return dedxMean;
}
