/**
 * @file
 * @author Christian Simon <csimon@physi.uni-heidelberg.de>
 * @since 2017-07-18
 */

#include "CbmTofFloodIlluminator.h"

#include "CbmTarget.h"

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

#include "TRandom3.h"
#include "TParticlePDG.h"
#include "TDatabasePDG.h"
#include "TGeoManager.h"
#include "TGeoPhysicalNode.h"

#include <boost/regex.hpp>
#include <boost/lexical_cast.hpp>


// ---------------------------------------------------------------------------
CbmTofFloodIlluminator::CbmTofFloodIlluminator() :
  FairGenerator(),
  fiPDGType(13),
  fdMomentum(10.),
  fiModuleType(0),
  fiModuleIndex(0),
  fiCounterIndex(0),
  fdWindowXLow(-1.),
  fdWindowXHigh(-1.),
  fdWindowYLow(-1.),
  fdWindowYHigh(-1.),
  fdCenterX(0.),
  fdCenterY(0.),
  fdIlluminationXLow(0.),
  fdIlluminationXHigh(0.),
  fdIlluminationYLow(0.),
  fdIlluminationYHigh(0.),
  fbIsInitialized(kFALSE),
  fTofNodePath(""),
  fCurrentNodePath(""),
  fiCurrentModuleType(-1),
  fiCurrentModuleIndex(-1),
  fiCurrentCounterIndex(-1),
  fIlluminationNode(NULL),
  fIlluminationPN(NULL),
  fiNEvents(0),
  fbXYPlaneVertex(kFALSE),
  fbOriginVertex(kFALSE),
  fbShiftWindowToYZPlane(kFALSE),
  fbSmearGausBeamOnTarget(kFALSE),
  fTarget(NULL)
{
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
CbmTofFloodIlluminator::~CbmTofFloodIlluminator()
{
  if(fIlluminationPN)
  {
    delete fIlluminationPN;
  }
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
void CbmTofFloodIlluminator::SetMomentum(Double_t dMomentum)
{
  if(0. >= dMomentum)
  {
    LOG(FATAL)<<Form("A non-positive particle momentum of %f GeV/c is not meaningful.", dMomentum)<<FairLogger::endl;
  }

  FairRunSim::Instance()->SetBeamMom(dMomentum);
  fdMomentum = dMomentum;
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
void CbmTofFloodIlluminator::SetCounter(Int_t iModuleType, Int_t iModuleIndex, Int_t iCounterIndex)
{
  fiModuleType = iModuleType;
  fiModuleIndex = iModuleIndex;
  fiCounterIndex = iCounterIndex;
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
void CbmTofFloodIlluminator::SetWindow(Double_t dXLow, Double_t dXHigh, Double_t dYLow, Double_t dYHigh, Double_t dCenterX, Double_t dCenterY)
{
  fdWindowXLow = dXLow;
  fdWindowXHigh = dXHigh;
  fdWindowYLow = dYLow;
  fdWindowYHigh = dYHigh;
  fdCenterX = dCenterX;
  fdCenterY = dCenterY;
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
Bool_t CbmTofFloodIlluminator::Init()
{
  // The complete TGeoManager to be used in the MC simulation is not available
  // yet when this method is called for the first time indirectly via
  // 'FairMCApplication::AddIons'. Thus, initialization of 'CbmTofFloodIllumination'
  // needs to be postponed to after 'FairRunSim::Init' has been called in the
  // macro. For this reason, upon the first call of 'Init' we just return 'kTRUE'
  // and register the method call with 'fbIsInitialized'.
  if(!fbIsInitialized)
  {
    fbIsInitialized = kTRUE;
    return kTRUE;
  }

  if(fbXYPlaneVertex && fbOriginVertex)
  {
    LOG(FATAL)<<"Cannot request XY plane and origin vertices simultaneously."<<FairLogger::endl;
  }

  // check for particle type
  TDatabasePDG* tPDGBase = TDatabasePDG::Instance();
  TParticlePDG* tParticle = tPDGBase->GetParticle(fiPDGType);

  if(!tParticle)
  {
    LOG(FATAL)<<Form("Particle of type %d not found in PDG database.", fiPDGType)<<FairLogger::endl;
  }

  if(!FindIlluminationNode())
  {
    LOG(FATAL)<<Form("Failed to extract module %d-%d-%d from TGeoManager.", fiModuleType, fiModuleIndex, fiCounterIndex)<<FairLogger::endl;
  }

  if(fbShiftWindowToYZPlane)
  {
    Double_t dLocalCounterCenter[3] = {0., 0., 0.};
    Double_t dGlobalCounterCenter[3] = {0., 0., 0.};

    fIlluminationPN->GetMatrix()->LocalToMaster(dLocalCounterCenter, dGlobalCounterCenter);

    // shift the local window center along the global X axis to compensate for the counter
    // center's shift w.r.t. the global YZ plane
    dGlobalCounterCenter[0] = -dGlobalCounterCenter[0];
    dGlobalCounterCenter[1] = 0.;
    dGlobalCounterCenter[2] = 0.;
    Double_t dLocalCenterShift[3] = {0., 0., 0.};

    fIlluminationPN->GetMatrix()->MasterToLocalVect(dGlobalCounterCenter, dLocalCenterShift);

    fdCenterX += dLocalCenterShift[0];
    fdCenterY += dLocalCenterShift[1];
  }

  TGeoVolume* tGlassPlateVolume = fIlluminationNode->GetVolume();
  TGeoShape* tGlassPlateShape = tGlassPlateVolume->GetShape();

  Double_t dAxisLow(0.);
  Double_t dAxisHigh(0.);

  Double_t dGlassPlateWidthX = tGlassPlateShape->GetAxisRange(1, dAxisLow, dAxisHigh);
  Double_t dGlassPlateWidthY = tGlassPlateShape->GetAxisRange(2, dAxisLow, dAxisHigh);

  Double_t dLocalCoordinates[3] = {fdCenterX, fdCenterY, 0.};

  if(!fIlluminationPN->GetShape()->Contains(dLocalCoordinates))
  {
    LOG(FATAL)<<Form("Illumination window center (%f,%f) is outside the local counter volume range.", fdCenterX, fdCenterY)<<FairLogger::endl;
  }

  // If ALL spatial extensions of the illumination window are set to/kept at '-1.',
  // the full counter surface will be irradiated, regardless of the irradiation center.
  // If ALL spatial extensions are positive, the counter will be irradiated about the
  // given irradiation center in the respective direction as specified by the
  // extension parameter but at most up to its volume boundary.
  // If any spatial extension of the illumination window is set negative or '0.', a
  // point-like illumination will be generated.
  if(-1. == fdWindowXLow && -1. == fdWindowXHigh && -1. == fdWindowYLow && -1. == fdWindowYHigh)
  {
    fdIlluminationXLow = -dGlassPlateWidthX/2.;
    fdIlluminationXHigh = dGlassPlateWidthX/2.;
    fdIlluminationYLow = -dGlassPlateWidthY/2.;
    fdIlluminationYHigh = dGlassPlateWidthY/2.;
  }
  else if(0. < fdWindowXLow && 0. < fdWindowXHigh && 0. < fdWindowYLow && 0. < fdWindowYHigh)
  {
    if(fdCenterX - fdWindowXLow < -dGlassPlateWidthX/2.)
    {
      fdIlluminationXLow = -dGlassPlateWidthX/2.;
    }
    else
    {
      fdIlluminationXLow = fdCenterX - fdWindowXLow;
    }

    if(fdCenterX + fdWindowXHigh > dGlassPlateWidthX/2.)
    {
      fdIlluminationXHigh = dGlassPlateWidthX/2.;
    }
    else
    {
      fdIlluminationXHigh = fdCenterX + fdWindowXHigh;
    }

    if(fdCenterY - fdWindowYLow < -dGlassPlateWidthY/2.)
    {
      fdIlluminationYLow = -dGlassPlateWidthY/2.;
    }
    else
    {
      fdIlluminationYLow = fdCenterY - fdWindowYLow;
    }

    if(fdCenterY + fdWindowYHigh > dGlassPlateWidthY/2.)
    {
      fdIlluminationYHigh = dGlassPlateWidthY/2.;
    }
    else
    {
      fdIlluminationYHigh = fdCenterY + fdWindowYHigh;
    }
  }

  LOG(INFO)<<Form("Illumination window centered at (%f,%f). Extension in X: [%f,%f]. Extension in Y: [%f,%f]",
                  fdCenterX, fdCenterY, fdIlluminationXLow, fdIlluminationXHigh, fdIlluminationYLow, fdIlluminationYHigh)
           <<FairLogger::endl;

  // Disable any modifications to be made to the primary vertex by 'FairPrimaryGenerator'
  FairPrimaryGenerator* tGenerator = FairRunSim::Instance()->GetPrimaryGenerator();
  tGenerator->SetTarget(0., 0.);
  tGenerator->SetBeam(0., 0., 0., 0.);
  tGenerator->SetBeamAngle(0., 0., 0., 0.);
  tGenerator->SetEventPlane(0., 0.);
  tGenerator->SmearVertexZ(kFALSE);
  tGenerator->SmearGausVertexZ(kFALSE);
  tGenerator->SmearVertexXY(kFALSE);
  tGenerator->SmearGausVertexXY(kFALSE);

  if(fbOriginVertex)
  {
    if(fbSmearGausBeamOnTarget)
    {
      if(!fTarget)
      {
        LOG(FATAL)<<"Pointer to 'CbmTarget' object not set despite requesting Gaussian beam smearing on target."<<FairLogger::endl;
      }

      Double_t dTargetX(0.);
      Double_t dTargetY(0.);
      Double_t dTargetZ(0.);

      fTarget->GetPosition(dTargetX, dTargetY, dTargetZ);

      if(0. != dTargetX || 0. != dTargetY || 0. != dTargetZ)
      {
        LOG(FATAL)<<"'CbmTarget' object is not positioned at the origin of coordinates (0, 0, 0)."<<FairLogger::endl;
      }

      Double_t dTargetR = fTarget->GetDiameter()/2.;

      // 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.
      tGenerator->SetBeam(0., 0., dTargetR/5., dTargetR/5.);
      tGenerator->SmearGausVertexXY(kTRUE);
    }
  }

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



// ---------------------------------------------------------------------------
Bool_t CbmTofFloodIlluminator::ReadEvent(FairPrimaryGenerator* primGen)
{
  FairMCEventHeader* tEventHeader = primGen->GetEvent();

  Double_t dLocalPositionVector[3] = {gRandom->Uniform(fdIlluminationXLow, fdIlluminationXHigh),
                                      gRandom->Uniform(fdIlluminationYLow, fdIlluminationYHigh),
                                      0.
                                     };

  Double_t dLocalMomentumVector[3] = {0., 0., fdMomentum};

  Double_t dGlobalPositionVector[3] = {0., 0., 0.};
  Double_t dGlobalMomentumVector[3] = {0., 0., 0.};

  fIlluminationPN->GetMatrix()->LocalToMaster(dLocalPositionVector, dGlobalPositionVector);

  if(fbOriginVertex)
  {
    // If 'fbSmearGausBeamOnTarget' is true, 'FairMCEventHeader::fX' and '::fY'
    // contain the Gaussian smearing of the primary vertex as provided by
    // 'FairPrimaryGenerator::GenerateEvent'.
    Double_t dVectorMagnitude = TMath::Sqrt( TMath::Power(dGlobalPositionVector[0] - tEventHeader->GetX(), 2.)
                                            +TMath::Power(dGlobalPositionVector[1] - tEventHeader->GetY(), 2.)
                                            +TMath::Power(dGlobalPositionVector[2], 2.));

    if(0. == dVectorMagnitude)
    {
      LOG(FATAL)<<"Illumination volume contains the origin. Requesting origin vertices is not meaningful."<<FairLogger::endl;
    }

    dGlobalMomentumVector[0] = (dGlobalPositionVector[0] - tEventHeader->GetX())*fdMomentum/dVectorMagnitude;
    dGlobalMomentumVector[1] = (dGlobalPositionVector[1] - tEventHeader->GetY())*fdMomentum/dVectorMagnitude;
    dGlobalMomentumVector[2] = dGlobalPositionVector[2]*fdMomentum/dVectorMagnitude;
  }
  else
  {
    fIlluminationPN->GetMatrix()->LocalToMasterVect(dLocalMomentumVector, dGlobalMomentumVector);
  }

  Double_t dGlobalTrackVertex[3] = {0., 0., 0.};

  // Extrapolate the particle track which is locally orthogonal to the counter plane
  // to the global XY coordinate plane upon request. If the local counter plane for whatever
  // reason happened to be parallel to the global XZ coordinate plane, this request could
  // not be met mathematically.
  if(fbXYPlaneVertex && 0. != dGlobalMomentumVector[2])
  {
    dGlobalTrackVertex[0] = dGlobalPositionVector[0] - dGlobalMomentumVector[0]/dGlobalMomentumVector[2]*dGlobalPositionVector[2];
    dGlobalTrackVertex[1] = dGlobalPositionVector[1] - dGlobalMomentumVector[1]/dGlobalMomentumVector[2]*dGlobalPositionVector[2];
    dGlobalTrackVertex[2] = 0.;
  }
  else if(fbOriginVertex)
  {
    dGlobalTrackVertex[0] = 0.;
    dGlobalTrackVertex[1] = 0.;
    dGlobalTrackVertex[2] = 0.;
  }
  else
  {
    dGlobalTrackVertex[0] = dGlobalPositionVector[0];
    dGlobalTrackVertex[1] = dGlobalPositionVector[1];
    dGlobalTrackVertex[2] = dGlobalPositionVector[2];
  }

  tEventHeader->SetEventID(fiNEvents); // Make sure that FairMCEventHeader::fEventId starts from 0!
  tEventHeader->SetB(0.);
  tEventHeader->SetNPrim(1);           // (re-)set in FairPrimaryGenerator::GenerateEvent
  tEventHeader->SetVertex(dGlobalTrackVertex[0] + tEventHeader->GetX(), dGlobalTrackVertex[1] + tEventHeader->GetY(), dGlobalTrackVertex[2]);
  tEventHeader->SetTime(0.);
  tEventHeader->SetRotX(0.);
  tEventHeader->SetRotY(0.);
  tEventHeader->SetRotZ(0.);
  tEventHeader->MarkSet(kTRUE);

  primGen->AddTrack(fiPDGType, dGlobalMomentumVector[0], dGlobalMomentumVector[1], dGlobalMomentumVector[2],
                               dGlobalTrackVertex[0], dGlobalTrackVertex[1], dGlobalTrackVertex[2]);

  fiNEvents++;

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



// ---------------------------------------------------------------------------
Bool_t CbmTofFloodIlluminator::FindIlluminationNode()
{
  if(!gGeoManager)
  {
    LOG(ERROR)<<"TGeoManager instance not found."<<FairLogger::endl;
    return kFALSE;
  }

  Bool_t bFoundTofNode(kFALSE);
  gGeoManager->CdTop();

  // Just in case a geometry file contains a TGeoManager with a 'tof' top node
  if(!TString(gGeoManager->GetCurrentNode()->GetName()).Contains("tof",TString::kIgnoreCase))
  {
    // Normally, the top node should be a 'cave' by which the 'tof' node is contained
    TObjArray* tNodes = gGeoManager->GetCurrentNode()->GetNodes();
    for(Int_t iNode = 0; iNode < tNodes->GetEntriesFast(); iNode++)
    {
      TGeoNode* tNode = dynamic_cast<TGeoNode*>(tNodes->At(iNode));
      if(TString(tNode->GetName()).Contains("tof",TString::kIgnoreCase))
      {
        gGeoManager->CdDown(iNode);
        bFoundTofNode = kTRUE;
        break;
      }
    }

    if(!bFoundTofNode)
    {
      LOG(ERROR)<<"Could not retrieve 'tof' node from TGeoManager."<<FairLogger::endl;
      return kFALSE;
    }
  }

  fTofNodePath = gGeoManager->GetPath();

  ExpandNode(gGeoManager->GetCurrentNode());

  gGeoManager->CdTop();

  if(!fIlluminationPN)
  {
    return kFALSE;
  }

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



// ---------------------------------------------------------------------------
void CbmTofFloodIlluminator::ExpandNode(TGeoNode* tMotherNode)
{
  TObjArray* tDaughterNodesList = tMotherNode->GetNodes();

  for(Int_t iNode = 0; iNode < tDaughterNodesList->GetEntriesFast(); iNode++)
  {
    TGeoNode* tDaughterNode = dynamic_cast<TGeoNode*>(tDaughterNodesList->At(iNode));

    // Extract the current module type and module index from the module node
    if(TString(tDaughterNode->GetName()).Contains("module"))
    {
      fiCurrentModuleType = -1;
      fiCurrentModuleIndex = -1;
      fiCurrentCounterIndex = -1;

      fCurrentNodePath = fTofNodePath + "/" + (TString)tDaughterNode->GetName();

      boost::regex rgx(".*_(\\d+)_.*");
      boost::cmatch match;
      if( boost::regex_search(tDaughterNode->GetName(), match, rgx) )
      {
        fiCurrentModuleType = boost::lexical_cast<Int_t>(match[1]);
      }

      fiCurrentModuleIndex = tDaughterNode->GetNumber();
    }
    // Extract the current counter index from the counter node and find the
    // central node in the counter volume (a glass plate OR a gap) to serve as
    // reference coordinate system
    else if(TString(tDaughterNode->GetName()).Contains("counter"))
    {
      fCurrentNodePath += "/" + (TString)tDaughterNode->GetName();

      fiCurrentCounterIndex = tDaughterNode->GetNumber();

      if(fiCurrentModuleType == fiModuleType &&
         fiCurrentModuleIndex == fiModuleIndex &&
         fiCurrentCounterIndex == fiCounterIndex)
      {
        // Here, we safely assume that every "counter" in the geometry contains 
        // a stack of glass plates and gaps. The two framing volumes of the stack
        // are glass plates. As the flood illumination is supposed to originate
        // in the framing glass plate and to point along the local Z axis we check here
        // whether the plate with the lowest node index is located in the negative local
        // Z range (DEFAULT) or the plate with the highest node index.
        Double_t dFirstGlassZPosition = (tDaughterNode->GetDaughter(0)->GetMatrix()->GetTranslation())[2];
        Double_t dLastGlassZPosition = (tDaughterNode->GetDaughter(tDaughterNode->GetNdaughters()-1)->GetMatrix()->GetTranslation())[2];

        if(dFirstGlassZPosition < dLastGlassZPosition)
        {
          fIlluminationNode = tDaughterNode->GetDaughter(0);
          fIlluminationPN = new TGeoPhysicalNode((fCurrentNodePath + TString::Format("/tof_glass_0")).Data());
        }
        else
        {
          fIlluminationNode = tDaughterNode->GetDaughter(tDaughterNode->GetNdaughters()-1);
          fIlluminationPN = new TGeoPhysicalNode((fCurrentNodePath + TString::Format("/tof_glass_%d",tDaughterNode->GetNdaughters()-1)).Data());
        }
      }
    }
    else
    {
      fCurrentNodePath += "/" + (TString)tDaughterNode->GetName();
    }

    // Expand nodes recursively
    if(tDaughterNode->GetNdaughters() > 0) { ExpandNode(tDaughterNode); }

    // Remove a node's name from the current node path upon completing a
    // recursion step
    fCurrentNodePath.ReplaceAll("/"+(TString)tDaughterNode->GetName(), "");
  }
}
// ---------------------------------------------------------------------------

ClassImp(CbmTofFloodIlluminator)