/**
 * @file
 * @author Christian Simon <csimon@physi.uni-heidelberg.de>
 * @since 2018-08-02
 */


#include "CbmTofFindPointType.h"

#include "CbmTofDigiTbPar.h"
#include "CbmTofCounter.h"
#include "CbmTofWall.h"
#include "CbmTofPoint.h"
#include "CbmMCTrack.h"
#include "CbmMCDataManager.h"
#include "CbmMCDataArray.h"
#include "CbmTofDef.h"

#include "FairRun.h"
#include "FairRootManager.h"
#include "FairRuntimeDb.h"
#include "FairLogger.h"

#include "TClonesArray.h"
#include "TGeoManager.h"



// ---------------------------------------------------------------------------
CbmTofFindPointType::CbmTofFindPointType() :
  FairTask("CbmTofFindPointType", 0),
  fMCTracks(NULL),
  fTofPointsTB(NULL),
  fDigiTbParSet(NULL),
  fPointClusterMap(),
  fReferencePointMap(),
  fbAlternativeBranchNames(kFALSE),
  fiFileIndex(0),
  fbConsiderPointClusters(kFALSE),
  fbConsiderRefTracks(kFALSE)
{
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
CbmTofFindPointType::~CbmTofFindPointType()
{
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
void CbmTofFindPointType::Exec(Option_t */*option*/)
{
  fPointClusterMap.clear();
  fReferencePointMap.clear();

  Int_t iTreeEntry = FairRootManager::Instance()->GetEntryNr();

  // Create point clusters and the corresponding maps
  if(fbConsiderPointClusters)
  {
    Double_t dLocalTrackStart[3] = {0., 0., 0.};
    Double_t dGlobalTrackStart[3] = {0., 0., 0.};
    Double_t dLocalMotherTrackStart[3] = {0., 0., 0.};
    Double_t dGlobalMotherTrackStart[3] = {0., 0., 0.};

    for(Int_t iPoint = 0; iPoint < fTofPointsTB->GetEntriesFast(); iPoint++)
    {
      CbmTofPoint* tPoint = dynamic_cast<CbmTofPoint*>(fTofPointsTB->At(iPoint));

      // Extract the file ID of the associated MC track from 'CbmTofPoint::fUniqueID'
      // before reusing this member variable to store the CbmTofPointType of
      // the MC point for further online analysis (the modified point object
      // is not written to disk).
      Int_t iFileID = tPoint->GetUniqueID();

      // Mark all ToF points in the tree entry as internal cluster points by default.
      tPoint->SetUniqueID(tof::point_Internal);

      Double_t dPointTime = tPoint->GetTime();

      Int_t iMCEvent = tPoint->GetEventID();
      Int_t iTrackID = tPoint->GetTrackID();
      Int_t iCounterID = tPoint->GetDetectorID();

      CbmTofCounter* tCounter = CbmTofWall::Instance()->GetCounter(iCounterID);
      auto & tCounterPointClusters = fPointClusterMap[iCounterID];
      CbmTofPointCluster* tPointCluster(NULL);

      // beam point or dark point (no secondaries available)
      if(-1 == iTrackID)
      {
        // Only consider beam points here as all dark points would be assigned
        // to a single 'CbmTofPointCluster' due to their (-1, -1) track ID pair.
        // To every counter at most one beam point can be assigned per event
        // which makes beam point clusters (always of size 1) distinguishable.
        // -> special treatment for dark points required later on!
        if(-1 != iMCEvent)
        {
          tPointCluster = &(((tCounterPointClusters.emplace(std::make_tuple(iFileID, iMCEvent, iTrackID), CbmTofPointCluster())).first)->second);
          tPointCluster->SetCharged(kTRUE);
          tPointCluster->AddPoint(fiFileIndex, iTreeEntry, iPoint, dPointTime, kTRUE);
        }
        else
        {
          // As no point clusters are created for dark points, their 'fUniqueID'
          // needs to be set here.
          tPoint->SetUniqueID(tof::point_Dark);
        }
      }
      else
      {
        CbmMCTrack* tTrack = dynamic_cast<CbmMCTrack*>(fMCTracks->Get(iFileID, iMCEvent, iTrackID));

        dGlobalTrackStart[0] = tTrack->GetStartX();
        dGlobalTrackStart[1] = tTrack->GetStartY();
        dGlobalTrackStart[2] = tTrack->GetStartZ();

        tCounter->MasterToLocalBox(dGlobalTrackStart, dLocalTrackStart);

        Int_t iMotherID = tPoint->GetTrackID();
        Int_t iKinshipDegree = 0;
        CbmMCTrack* tMotherTrack = NULL;

        while(-1 != iMotherID)
        {
          tMotherTrack = dynamic_cast<CbmMCTrack*>(fMCTracks->Get(iFileID, iMCEvent, iMotherID));

          dGlobalMotherTrackStart[0] = tMotherTrack->GetStartX();
          dGlobalMotherTrackStart[1] = tMotherTrack->GetStartY();
          dGlobalMotherTrackStart[2] = tMotherTrack->GetStartZ();

          tCounter->MasterToLocalBox(dGlobalMotherTrackStart, dLocalMotherTrackStart);

          // lowest-ranked mother particle created outside the counter box volume
          if(!tCounter->IsContainedInBox(dLocalMotherTrackStart))
          {
            std::tuple<Int_t, Int_t, Int_t> tMotherTag(iFileID, iMCEvent, iMotherID);

            auto itTemp = tCounterPointClusters.find(tMotherTag);

            if(itTemp == tCounterPointClusters.end())
            {
              tPointCluster = &tCounterPointClusters[tMotherTag];

              if(!tMotherTrack->GetCharge())
              {
                tPointCluster->SetCharged(kFALSE);
              }

              tPointCluster->AddPoint(fiFileIndex, iTreeEntry, iPoint, dPointTime, 0 == iKinshipDegree);
            }
            else
            {
              tPointCluster = &tCounterPointClusters.at(tMotherTag);

              if(tCounter->IsContainedInBox(dLocalTrackStart))
              {
                tPointCluster->AddPoint(fiFileIndex, iTreeEntry, iPoint, dPointTime, 0 == iKinshipDegree);
              }
            }

            break;
          }

          iMotherID = tMotherTrack->GetMotherId();
          iKinshipDegree++;
        }
      }
    }


    for(auto const & itCounter : fPointClusterMap)
    {
      for(auto const & itCluster: itCounter.second)
      {
        const CbmTofPointCluster* tPointCluster = &itCluster.second;

        // TODO: add an option to select direct points only
        fReferencePointMap.emplace(tPointCluster->GetReferencePoint(), tPointCluster->IsCharged() ? tof::point_ExternalCharged : tof::point_ExternalNeutral);
      }
    }


    for(Int_t iPoint = 0; iPoint < fTofPointsTB->GetEntriesFast(); iPoint++)
    {
      CbmTofPoint* tPoint = dynamic_cast<CbmTofPoint*>(fTofPointsTB->At(iPoint));

      auto itTemp = fReferencePointMap.find(std::make_tuple(fiFileIndex, iTreeEntry, iPoint));

      if(itTemp != fReferencePointMap.end())
      {
        tPoint->SetUniqueID(itTemp->second);
      }
    }

  }
  // Find reference track points
  else if(fbConsiderRefTracks)
  {
    TGeoNode* tNode(NULL);
    TGeoMedium* tMedium(NULL);
    TGeoMaterial* tMaterial(NULL);

    const char* cMaterialName;

    for(Int_t iPoint = 0; iPoint < fTofPointsTB->GetEntriesFast(); iPoint++)
    {
      CbmTofPoint* tPoint = dynamic_cast<CbmTofPoint*>(fTofPointsTB->At(iPoint));

      // Extract the file ID of the associated MC track from 'CbmTofPoint::fUniqueID'
      // before reusing this member variable to store the CbmTofPointType of
      // the MC point for further online analysis (the modified point object
      // is not written to disk).
      Int_t iFileID = tPoint->GetUniqueID();
      Int_t iMCEvent = tPoint->GetEventID();
      Int_t iTrackID = tPoint->GetTrackID();

      // Mark all ToF points in the TS as internal cluster points by default.
      tPoint->SetUniqueID(tof::point_Internal);

      // MC track object cannot be retrieved.
      if(-1 == iTrackID)
      {
        // Mark beam points for hit building.
        if(-1 != iMCEvent)
        {
          fReferencePointMap.emplace(std::make_tuple(fiFileIndex, iTreeEntry, iPoint), tof::point_RefTrack);
          tPoint->SetUniqueID(tof::point_RefTrack);
        }
      }
      // MC track object can be retrieved.
      else
      {
        CbmMCTrack* tTrack = dynamic_cast<CbmMCTrack*>(fMCTracks->Get(iFileID, iMCEvent, iTrackID));

        tNode = gGeoManager->FindNode(tTrack->GetStartX(), tTrack->GetStartY(), tTrack->GetStartZ());
        tMedium = tNode->GetMedium();
        tMaterial = tMedium->GetMaterial();

        tof::GetMaterialName(tMaterial->GetName(), cMaterialName);

        if(0 == std::strcmp("target", cMaterialName) && 3 <= tTrack->GetNPoints(kTof))
        {
          fReferencePointMap.emplace(std::make_tuple(fiFileIndex, iTreeEntry, iPoint), tof::point_RefTrack);
          tPoint->SetUniqueID(tof::point_RefTrack);
        }
      }
    }
  }
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
InitStatus CbmTofFindPointType::Init()
{
  if(fbConsiderPointClusters && fbConsiderRefTracks)
  {
    LOG(FATAL)<<"Cannot request point cluster and reference track selection simultaneously."<<FairLogger::endl;
  }

  if(!(CbmTofWall::Instance()->Initialize(fDigiTbParSet)))
  {
    return kFATAL;
  }

  // for future use
  for(auto const & itCounter : CbmTofWall::Instance()->GetCounterMap())
  {
    (itCounter.second)->SetFileIndex(fiFileIndex);
  }

	// --- Get FairRootManager instance
  FairRootManager* tRootManager = FairRootManager::Instance();
  if(!tRootManager)
  {
    LOG(ERROR)<<"FairRootManager not found."<<FairLogger::endl;
    return kFATAL;
  }


  TString tPointBranchName;

  if(fbAlternativeBranchNames)
  {
    tPointBranchName = "ATofPointTB";
  }
  else
  {
    tPointBranchName = "TofPointTB";
  }

  // --- Get input array (CbmTofPoint)
  fTofPointsTB = dynamic_cast<TClonesArray*>(tRootManager->GetObject(tPointBranchName));
  if(!fTofPointsTB)
  {
    LOG(ERROR)<<Form("Could not retrieve branch %s from FairRootManager.", tPointBranchName.Data())<<FairLogger::endl;
    return kFATAL;
  }

  CbmMCDataManager* tMCManager = dynamic_cast<CbmMCDataManager*>(tRootManager->GetObject("MCDataManager"));
  if(!tMCManager)
  {
    LOG(ERROR)<<"Could not retrieve CbmMCDataManager from FairRootManager."<<FairLogger::endl;
    return kFATAL;
  }

  fMCTracks = tMCManager->InitBranch("MCTrack");
  if(!fMCTracks)
  {
    LOG(ERROR)<<"Could not retrieve branch MCTrack from CbmMCDataManager."<<FairLogger::endl;
    return kFATAL;
  }


  return kSUCCESS;
}
// ---------------------------------------------------------------------------



// ---------------------------------------------------------------------------
void CbmTofFindPointType::SetParContainers()
{
  FairRun* tRun = FairRun::Instance();
  if(!tRun)
  {
    LOG(FATAL)<<Form("FairRun instance not found.")<<FairLogger::endl;
  }

  FairRuntimeDb* tDataBase = tRun->GetRuntimeDb();
  if(!tDataBase)
  {
    LOG(FATAL)<<Form("FairRuntimeDb instance not found.")<<FairLogger::endl;
  }

  // request initialization of parameter containers
  // The actual parameter container initialization happens only upon calling
  // FairRuntimeDb::initContainers in the FairRun instance
  fDigiTbParSet = dynamic_cast<CbmTofDigiTbPar*>(tDataBase->getContainer("CbmTofDigiTbPar"));
}
// ---------------------------------------------------------------------------


ClassImp(CbmTofFindPointType)