/**
 * @file
 * @author Dmytro Kresan <d.kresan@gsi.de>
 * @author Christian Simon <csimon@physi.uni-heidelberg.de>
 * @since 2006-06-04/2018-11-23
 */

#include "CbmTofUnigenGenerator.h"

#include "FairMCEventHeader.h"

#include "FairPrimaryGenerator.h"
#include "FairLogger.h"
#include "FairRunSim.h"
#include "FairRootManager.h"

#include "URun.h"
#include "UEvent.h"
#include "UParticle.h"
#include "CbmTarget.h"

#include "TFile.h"
#include "TChain.h"
#include "TRandom.h"
#include "TF1.h"
#include "TVectorD.h"
#include "TRegexp.h"
#include "TSystemDirectory.h"
#include "TSystem.h"

#include <iostream>
#include <algorithm>
using namespace std;


// ------------------------------------------------------------------------
CbmTofUnigenGenerator::CbmTofUnigenGenerator()
  : FairGenerator(),
    fEvents(0),
    fMaxEvents(0),
    fInputFile(NULL),
    fFileName(""),
    fInChain(NULL),
    fInputChainList(),
    fEvent(NULL),
    fRun(NULL),
    fCM(kFALSE),
    fBetaCM(0.),
    fGammaCM(1.),
    fBetaProj(0.),
    fTarget(NULL),
    fTargetZPosition(0.),
    fTargetZLength(0.),
    fBeamSigmaX(0.),
    fBeamSigmaY(0.),
    fCollZProb(-1.),
    fSmearVertexXY(kFALSE),
    fSmearGausVertexXY(kFALSE),
    fSmearVertexZ(kFALSE),
    fSmearExpVertexZ(kFALSE),
    fSampleZPosition(NULL),
    fAutoCorrectBeam(kFALSE),
    fAllGeantinos(kFALSE)
{
}
// ------------------------------------------------------------------------


// ------------------------------------------------------------------------
CbmTofUnigenGenerator::CbmTofUnigenGenerator(TString fileName)
  : FairGenerator(),
    fEvents(0),
    fMaxEvents(0),
    fInputFile(NULL),
    fFileName(fileName),
    fInChain(NULL),
    fInputChainList(),
    fEvent(NULL),
    fRun(NULL),
    fCM(kFALSE),
    fBetaCM(0.),
    fGammaCM(0.),
    fBetaProj(0.),
    fTarget(NULL),
    fTargetZPosition(0.),
    fTargetZLength(0.),
    fBeamSigmaX(0.),
    fBeamSigmaY(0.),
    fCollZProb(-1.),
    fSmearVertexXY(kFALSE),
    fSmearGausVertexXY(kFALSE),
    fSmearVertexZ(kFALSE),
    fSmearExpVertexZ(kFALSE),
    fSampleZPosition(NULL),
    fAutoCorrectBeam(kFALSE),
    fAllGeantinos(kFALSE)
{
  LOG(INFO) << "CbmTofUnigenGenerator: Opening input file "
            << fFileName.Data() << FairLogger::endl;
  fInputFile = new TFile(fFileName);
  if (NULL == fInputFile) {
    LOG(FATAL) << "CbmTofUnigenGenerator: Cannot open input file." << FairLogger::endl;
  }

  // Get run description
  fRun = (URun*) fInputFile->Get("run");
  if(NULL==fRun) {
    LOG(FATAL) << "CbmTofUnigenGenerator: No run description in input file." << FairLogger::endl;
  }
  fCM = kFALSE;
  LOG(DEBUG) << "Target Momentum: " << fRun->GetPTarg() << FairLogger::endl;
  if(TMath::Abs(fRun->GetPTarg()) > 0.001) { fCM = kTRUE; }
  fBetaCM = 0.;
  fGammaCM = 1.;
  if(fCM) {
    LOG(INFO) << "Input data is in CM frame" << FairLogger::endl;
    Double_t elab = (TMath::Power(fRun->GetNNSqrtS(),2)-
                     2*TMath::Power(0.938271998,2))/(2*0.938271998);
    Double_t plab = TMath::Sqrt(elab*elab - TMath::Power(0.938271998,2));
    LOG(INFO) << "CbmTofUnigenGenerator: Plab = " << plab << " AGeV" << FairLogger::endl;
    fBetaCM = plab / (elab + 0.938271998);
    fGammaCM = 1./TMath::Sqrt(1. - fBetaCM*fBetaCM);
    fBetaProj = plab/elab;
    LOG(INFO) << "CbmTofUnigenGenerator: boost parameters: "
              << "betaCM = "    << fBetaCM
              << ", gammaCM = " << fGammaCM << FairLogger::endl;
    // store the beam momentum in FairRunSim, it will be transfered to FairBaseParSet in FairRunSim::Init
    FairRunSim::Instance()->SetBeamMom(plab);
  } else {
    LOG(INFO) << "Input data is in LAB frame" << FairLogger::endl;
    fBetaProj = fRun->GetPProj()/TMath::Sqrt(TMath::Power(fRun->GetPProj(), 2.) + TMath::Power(0.938271998, 2.));
    FairRunSim::Instance()->SetBeamMom(fRun->GetPProj());
  }

  // Add input file name passed to constructor as first element to 'fInChain'
  AddFile(fFileName);
}
// ------------------------------------------------------------------------


// ------------------------------------------------------------------------
CbmTofUnigenGenerator::~CbmTofUnigenGenerator()
{
  if(fSampleZPosition) delete fSampleZPosition;

  CloseInputChain();
}
// ------------------------------------------------------------------------


// ------------------------------------------------------------------------
Bool_t CbmTofUnigenGenerator::Init()
{
  if(FairRunSim::Instance())
  {
    FairPrimaryGenerator* tGenerator = FairRunSim::Instance()->GetPrimaryGenerator();

    if(tGenerator)
    {
      // A 2D beam distribution with a single non-zero dimension does not make
      // too much sense.
      if(0. == fBeamSigmaX || 0. == fBeamSigmaY)
      {
        fBeamSigmaX = 0.;
        fBeamSigmaY = 0.;
      }

      if(fTarget)
      {
        Double_t dTargetX(0.);
        Double_t dTargetY(0.);

        Double_t dCollisionProbabilityXY(1.);
        Double_t dCollisionProbabilityZ(1.);

        // It is impossible to hit a point-like target.
        if(0. == fTarget->GetDiameter())
        {
          dCollisionProbabilityXY *= 0.;
        }

        // use target nuclear charge from input file
        fTarget->SetCharge(fRun->GetZTarg());

        // do not allow for a tilted target or beam
        fTarget->SetRotation(0.);
        tGenerator->SetBeamAngle(0., 0., 0., 0.);

        // disable unphysical Gaussian z vertex smearing
        tGenerator->SmearGausVertexZ(kFALSE);

        // Calculate the geometric mean free path of the beam particles from the input file
        // traversing the target medium from the input file.
        Double_t dNuclearR0 = 1.21e-13; // in cm

        Double_t dMeanFreePath = fTarget->GetMolarMass(fRun->GetZTarg())
                                /fTarget->GetStandardDensity(fRun->GetZTarg())
                                /TMath::Na()/TMath::Pi()
                                /TMath::Power(dNuclearR0, 2.)
                                /TMath::Power(std::cbrt(fRun->GetAProj()) + std::cbrt(fRun->GetATarg()), 2.);

        LOG(INFO) << "CbmTofUnigenGenerator: mean free path = " << dMeanFreePath << " cm" << FairLogger::endl;

        if(0. <= fCollZProb)
        {
          fTargetZLength = -TMath::Log(1. - fCollZProb)*dMeanFreePath;
          fTarget->SetThickness(fTargetZLength);
        }
        else
        {
          fTargetZLength = fTarget->GetThickness();
        }

        fTarget->GetPosition(dTargetX, dTargetY, fTargetZPosition);
        tGenerator->SetTarget(fTargetZPosition, fTargetZLength);

        dCollisionProbabilityZ *= 1. - TMath::Exp(-fTargetZLength/dMeanFreePath);


        tGenerator->SmearVertexXY(fSmearVertexXY);
        // The uniform beam profile implemented in FairPrimaryGenerator has
        // a rectangular shape. If the half-diagonal of the beam profile exceeds
        // the radius of the CbmTarget tube some share of beam particles is not
        // propagated through the target medium. In this case the beam profile
        // is auto-adjusted accordingly.
        if(fSmearVertexXY)
        {
          Double_t dBeamDiagonal = TMath::Sqrt(TMath::Power(fBeamSigmaX, 2.) + TMath::Power(fBeamSigmaY, 2.));
          Double_t dTargetDiameter = fTarget->GetDiameter();

          if(fAutoCorrectBeam && (dBeamDiagonal > dTargetDiameter))
          {
            LOG(WARNING) << "Rectangular beam profile is larger than circular target profile!" << FairLogger::endl;
            LOG(INFO) << "The beam profile is scaled down to fit the target profile." << FairLogger::endl;

            Double_t dProfileDownscale = dTargetDiameter/dBeamDiagonal;

            fBeamSigmaX *= dProfileDownscale;
            fBeamSigmaY *= dProfileDownscale;
          }

          Double_t dBeamProfile = fBeamSigmaX*fBeamSigmaY;
          Double_t dTargetProfile = TMath::Pi()/4.*dTargetDiameter*dTargetDiameter;
          Double_t dTargetOverlap(0.);

          // If the rectangular beam profile exceeds the circular target profile,
          // circular segment areas need to be calculated to compute the interaction
          // probability in xy. Else, all beam particles hit the target.
          if(dBeamDiagonal > dTargetDiameter)
          {
            if(fBeamSigmaX < dTargetDiameter && fBeamSigmaY >= dTargetDiameter)
            {
              dTargetOverlap = 0.5*( dTargetDiameter*dTargetDiameter*TMath::ACos(fBeamSigmaX/dTargetDiameter)
                                    -fBeamSigmaX*TMath::Sqrt(dTargetDiameter*dTargetDiameter - fBeamSigmaX*fBeamSigmaX));
            }
            else if(fBeamSigmaY < dTargetDiameter && fBeamSigmaX >= dTargetDiameter)
            {
              dTargetOverlap = 0.5*( dTargetDiameter*dTargetDiameter*TMath::ACos(fBeamSigmaY/dTargetDiameter)
                                    -fBeamSigmaY*TMath::Sqrt(dTargetDiameter*dTargetDiameter - fBeamSigmaY*fBeamSigmaY));
            }
            else if(fBeamSigmaX < dTargetDiameter && fBeamSigmaY < dTargetDiameter)
            {
              dTargetOverlap = 0.5*( dTargetDiameter*dTargetDiameter*TMath::ACos(fBeamSigmaX/dTargetDiameter)
                                    -fBeamSigmaX*TMath::Sqrt(dTargetDiameter*dTargetDiameter - fBeamSigmaX*fBeamSigmaX))
                              +0.5*( dTargetDiameter*dTargetDiameter*TMath::ACos(fBeamSigmaY/dTargetDiameter)
                                    -fBeamSigmaY*TMath::Sqrt(dTargetDiameter*dTargetDiameter - fBeamSigmaY*fBeamSigmaY));
            }

            dCollisionProbabilityXY *= (dTargetProfile - dTargetOverlap)/dBeamProfile;
          }

          LOG(INFO) << "CbmTofUnigenGenerator: xy interaction probability (rectangular beam) = " << dCollisionProbabilityXY << FairLogger::endl;
        }

        tGenerator->SmearGausVertexXY(fSmearGausVertexXY);
        // If the symmetric [-5s,+5s] interval (probability to lie outside
        // equals 1/1,744,278) of the marginal 1D Gaussian distribution of
        // any beam profile coordinate exceeds the target profile some
        // (possibly relevant) share of beam particles is not propagated through
        // the target medium. In this case the beam profile is auto-adjusted
        // accordingly.
        if(fSmearGausVertexXY)
        {
          Double_t dTargetRadius = fTarget->GetDiameter()/2.;

          if(fAutoCorrectBeam && (5.*TMath::Max(fBeamSigmaX, fBeamSigmaY) > dTargetRadius))
          {
            LOG(WARNING) << "Gaussian beam profile is larger than circular target profile!" << FairLogger::endl;
            LOG(INFO) << "The beam profile is scaled down to fit the target profile." << FairLogger::endl;

            Double_t dProfileDownscale = dTargetRadius/5./TMath::Max(fBeamSigmaX, fBeamSigmaY);

            fBeamSigmaX *= dProfileDownscale;
            fBeamSigmaY *= dProfileDownscale;
          }

          // The interaction probability between beam and target in xy here
          // corresponds to the 2D Gaussian beam distribution integrated over
          // the circular target profile.
          if(0. < fBeamSigmaX && 0. < fBeamSigmaY)
          {
            Int_t iNPoints = 1000;

            // For numerical simplification, one dimension of the 2D integral has
            // already been integrated out.
            TF1* tBeamIntegral = new TF1("tBeamIntegral","TMath::Exp(-0.5*[0]*[0]*x*x/([1]*[1]))*TMath::Erf([0]*TMath::Sqrt((1.-x*x)/(2.*[2]*[2])))", -1., 1.);
            tBeamIntegral->SetParameters(dTargetRadius, fBeamSigmaX, fBeamSigmaY);

            Double_t* dXi = new Double_t[iNPoints];
            Double_t* dWi = new Double_t[iNPoints];
            tBeamIntegral->CalcGaussLegendreSamplingPoints(iNPoints, dXi, dWi, 1e-15);

            dCollisionProbabilityXY *= dTargetRadius/TMath::Sqrt(TMath::TwoPi()*fBeamSigmaX*fBeamSigmaX)*tBeamIntegral->IntegralFast(iNPoints, dXi, dWi, -1., 1.);

            delete[] dXi;
            delete[] dWi;
            delete tBeamIntegral;
          }

          LOG(INFO) << "CbmTofUnigenGenerator: xy interaction probability (Gaussian beam) = " << dCollisionProbabilityXY << FairLogger::endl;
        }

        LOG(INFO) << "CbmTofUnigenGenerator: z interaction probability = " << dCollisionProbabilityZ << FairLogger::endl;
        LOG(INFO) << "CbmTofUnigenGenerator: total interaction probability = " << dCollisionProbabilityXY*dCollisionProbabilityZ << FairLogger::endl;

        // position the beam in the center of the target after cross-checking its profile
        tGenerator->SetBeam(dTargetX, dTargetY, fBeamSigmaX, fBeamSigmaY);
        LOG(INFO) << "CbmTofUnigenGenerator: beam width in x = " << fBeamSigmaX << " cm" << FairLogger::endl;
        LOG(INFO) << "CbmTofUnigenGenerator: beam width in y = " << fBeamSigmaY << " cm" << FairLogger::endl;

        tGenerator->SmearVertexZ(fSmearVertexZ);
        // Use the geometric mean free path of the beam in the target as parameter
        // of an exponential distribution describing the probability to collide
        // inside the target at [z,z+dz]. In CbmTofUnigenGenerator::ReadEvent, the
        // z vertex of each collision is then sampled from this distribution.
        if(fSmearExpVertexZ)
        {
          if(0. < fTargetZLength)
          {
            fSampleZPosition = new TF1("fSampleZPosition","1./[0]*TMath::Exp(-x/[0])", 0., fTargetZLength);
            fSampleZPosition->SetParameter(0, dMeanFreePath);
            fSampleZPosition->SetNpx(100000); // maybe not necessary (default: 100)
            fSampleZPosition->GetRandom();
          }
        }
/*
        // Write the total interaction probability to FairRootManager::fOutFile
        // to make it available during digitization when events need to be
        // positioned on a continuous time axis according to beam intensity AND
        // interaction probability (= event rate).
        TVectorD dInteractionProb(1);
        dInteractionProb[0] = dCollisionProbabilityXY*dCollisionProbabilityZ;

        TDirectory* tOldDirectory = gDirectory;
        FairRootManager::Instance()->GetOutFile()->cd();
        dInteractionProb.Write("InteractionProb");
        tOldDirectory->cd();
*/

        fTarget->Print();
      }
      else
      {
        // A valid CbmTarget is required to sample the z position of the collision
        // according to a realistic exponential distribution parametrized by the
        // geometric mean free path of the beam particles in the target medium.
        fSmearExpVertexZ = kFALSE;
      }
    }
  }

  delete fRun;

  CloseInputFile();

  if(!fInChain)
  {
    fInChain = new TChain("events");
    LOG(INFO) << "CbmTofUnigenGenerator::Init() chain created" << FairLogger::endl;

    // Add all input files to the input chain
    for(auto const & tFileName : fInputChainList)
    {
      TFile* inputFile = new TFile(tFileName);
      if(inputFile->IsZombie())
      {
	      LOG(FATAL) << "Error opening the file " << tFileName.Data() << " which should be added to the input chain." << FairLogger::endl;
      }
      else
      {
	      LOG(INFO) << "Adding file " << tFileName.Data() << " to the input chain." << FairLogger::endl;
      }

      inputFile->Close();
      delete inputFile;

      fInChain->Add(tFileName);
    }

    fInChain->SetBranchAddress("event", &fEvent);

    fMaxEvents = fInChain->GetEntries();
  }

  return kTRUE;
}
// ------------------------------------------------------------------------

// ------------------------------------------------------------------------
Bool_t CbmTofUnigenGenerator::ReadEvent(FairPrimaryGenerator* primGen)
{
  // Check for input file
  if(NULL == fInChain) {
    LOG(ERROR) << "CbmTofUnigenGenerator: Input chain is not available!" << FairLogger::endl;
    return kFALSE;
  }

  // If end of input chain is reached : close it and abort run
  if(fEvents >= fMaxEvents) {
    LOG(INFO) << "CbmTofUnigenGenerator: End of input chain reached" << FairLogger::endl;
    CloseInputChain();
    return kFALSE;
  }

  // Read entry
  fInChain->GetEntry(fEvents);

  LOG(INFO) << "CbmTofUnigenGenerator: Event " << fEvent->GetEventNr()
            << ", multiplicity " << fEvent->GetNpa() << FairLogger::endl;

  UParticle* particle;
  Double_t pz;
  Double_t pz1;
  Double_t vz(0.);
  Double_t tv(0.);

  // Set event id and impact parameter in MCEvent if not yet done
  FairMCEventHeader* event = primGen->GetEvent();
  if(event)
  {
    // FIXME Are there cases in which we use several FairGenerator instances in
    //       parallel? Why do we check FairMCEventHeader::IsSet here?
    if(!event->IsSet())
    {
      event->SetEventID(fEvents); // Make sure that FairMCEventHeader::fEventId starts from 0!
      event->SetB(fEvent->GetB());
      event->SetNPrim(fEvent->GetNpa()); // (re-)set in FairPrimaryGenerator::GenerateEvent
      event->MarkSet(kTRUE);
    }

    if(fSmearExpVertexZ)
    {
      Double_t dCollisionInTarget(0.);
      if(fSampleZPosition)
      {
        dCollisionInTarget = fSampleZPosition->GetRandom();
      }
      vz = dCollisionInTarget - fTargetZLength/2. + fTargetZPosition; // in cm
      tv = dCollisionInTarget/fBetaProj/TMath::Ccgs(); // in s relative to target z start

      event->SetVertex(event->GetX(), event->GetY(), vz);
    }
    else if(fSmearVertexZ)
    {
      // FairPrimaryGenerator samples an absolute collision z position.
      Double_t dCollisionInTarget = event->GetZ() - fTargetZPosition + fTargetZLength/2.;
      tv = dCollisionInTarget/fBetaProj/TMath::Ccgs(); // in s relative to target z start
    }
    
    event->SetTime(tv);
  }

  // Loop over tracks in the current event
  for (Int_t itrack = 0; itrack < fEvent->GetNpa(); itrack++) {
    // Get particle
    particle = fEvent->GetParticle(itrack);
    // Boost
    pz = particle->Pz();
    if(fCM) {
      pz1 = fGammaCM*(pz + fBetaCM*particle->E());
    } else {
      pz1 = pz;
    }
    Double_t px = particle->Px();
    Double_t py = particle->Py();

    if(fAllGeantinos)
    {
      particle->SetPdg(0);
    }

    // Give track to PrimaryGenerator
    primGen->AddTrack(particle->GetPdg(),
                      px,
                      py,
                      pz1,
                      0., 0., vz,
                      -1, kTRUE,
                      -9e9, tv,
                      0.);
  }

  fEvents += 1;

  return kTRUE;
}
// -------------------------------------------------------------------------

// -----   Public method AddFile   -----------------------------------------
void CbmTofUnigenGenerator::AddFile(const TString& FileName)
{
  if(find(fInputChainList.begin(), fInputChainList.end(), FileName) == fInputChainList.end())
  {
    // File name does not exist in list of file names
    fInputChainList.push_back(FileName);
  }
}
// -------------------------------------------------------------------------

// -----   Public method AddPath   -----------------------------------------
void CbmTofUnigenGenerator::AddPath(const TString& directoryName, const TString& fileNameWildCard)
{
  FileStat_t tFileStat;
  if(1 == gSystem->GetPathInfo(directoryName.Data(), tFileStat))
  {
    LOG(FATAL) << TString::Format("\nInput data file directory %s does not "
                                   "exist.",
                                   directoryName.Data()
                                  ).Data()
               << FairLogger::endl;
  }

  TRegexp* tRegexp = new TRegexp(fileNameWildCard.Data(), kTRUE);

  TSystemDirectory* tSystemDirectory = 
         new TSystemDirectory("dir", directoryName.Data());

  TString tTargetDirectoryName(directoryName);

  if(!tTargetDirectoryName.EndsWith("/"))
  {
      tTargetDirectoryName += "/";
  }

  TList* tList = tSystemDirectory->GetListOfFiles();
  tList->Sort();

  TIterator* tIter = tList->MakeIterator();
  TSystemFile* tSystemFile;

  TString tFileName;

  while(NULL != (tSystemFile = (TSystemFile*)tIter->Next()))
  {
    tFileName = tSystemFile->GetName();

    if(tFileName.Contains(*tRegexp))
    {
      tFileName = tTargetDirectoryName + tFileName;

      AddFile(tFileName);
    }
  }

  tList->Delete();
}
// -------------------------------------------------------------------------

// -----   Public method SetCollisionProbabilityZ   ------------------------
void CbmTofUnigenGenerator::SetCollisionProbabilityZ(Double_t collProbZ)
{
  if(0. <= collProbZ && collProbZ < 1.)
  {
    fCollZProb = collProbZ;
    fAutoCorrectBeam = kTRUE;
  }
  else
  {
    LOG(ERROR) << "CbmTofUnigenGenerator: collision probability in z "
              << collProbZ << " out of [0,1) range!" << FairLogger::endl;
  }
}
// -------------------------------------------------------------------------

// -----   Public method SmearVertexXY   -----------------------------------
void CbmTofUnigenGenerator::SmearVertexXY(Bool_t flag) {
  fSmearVertexXY = flag;
  if (fSmearVertexXY && fSmearGausVertexXY) {
    fSmearGausVertexXY = kFALSE;
  }
}
// -------------------------------------------------------------------------

// -----   Public method SmearGausVertexXY   -------------------------------
void CbmTofUnigenGenerator::SmearGausVertexXY(Bool_t flag) {
  fSmearGausVertexXY = flag;
  if (fSmearGausVertexXY && fSmearVertexXY) {
    fSmearVertexXY = kFALSE;
  }
}
// -------------------------------------------------------------------------

// -----   Public method SmearVertexZ   ------------------------------------
void CbmTofUnigenGenerator::SmearVertexZ(Bool_t flag) {
  fSmearVertexZ = flag;
  if (fSmearVertexZ && fSmearExpVertexZ) {
    fSmearExpVertexZ = kFALSE;
  }
}
// -------------------------------------------------------------------------

// -----   Public method SmearExpVertexZ   ---------------------------------
void CbmTofUnigenGenerator::SmearExpVertexZ(Bool_t flag) {
  fSmearExpVertexZ = flag;
  if (fSmearExpVertexZ && fSmearVertexZ) {
    fSmearVertexZ = kFALSE;
  }
}
// -------------------------------------------------------------------------

// ------------------------------------------------------------------------
void CbmTofUnigenGenerator::CloseInputFile()
{
  if(NULL != fInputFile) {
    LOG(INFO) << "CbmTofUnigenGenerator: Closing input file "
              << fFileName.Data() << FairLogger::endl;
    fInputFile->Close();
    delete fInputFile;
    fInputFile = NULL;
  }
}
// ------------------------------------------------------------------------

// ------------------------------------------------------------------------
void CbmTofUnigenGenerator::CloseInputChain()
{
  if(NULL != fInChain) {
    LOG(INFO) << "CbmTofUnigenGenerator: Closing input chain "
              << FairLogger::endl;
    fInChain->Reset();
  }
}
// ------------------------------------------------------------------------


ClassImp(CbmTofUnigenGenerator);