// -------------------------------------------------------------------------
// -----                       FairStack header file                    -----
// -----           Created 10/08/04  by D. Bertini                     -----
// -------------------------------------------------------------------------


/** FairStack.h
 *@author D.Bertini <d.bertini@gsi.de>
 * Generic MC stack class
 **/


#ifndef FAIRGENERICSTACK_H
#define FAIRGENERICSTACK_H


#include "FairDetector.h"
#include "FairMCPoint.h"

#include "TClonesArray.h"
#include "TVirtualMCStack.h"
#include "TRefArray.h"

#include <map>
#include <stack>
#include <list>

class TParticle;


class FairGenericStack : public TVirtualMCStack
{

  public:

    /** Default constructor  **/
    FairGenericStack();


    /** Destructor with estimated array size  **/
    FairGenericStack(Int_t size);


    /** Destructor  **/
    virtual ~FairGenericStack();

    /** Virtual method PushTrack.
     ** Add a TParticle to the stack.
     *@param toBeDone  Flag for tracking
     *@param parentID  Index of mother particle
     *@param pdgCode   Particle type (PDG encoding)
     *@param px,py,pz  Momentum components at start vertex [GeV]
     *@param e         Total energy at start vertex [GeV]
     *@param vx,vy,vz  Coordinates of start vertex [cm]
     *@param time      Start time of track [s]
     *@param polx,poly,polz Polarisation vector
     *@param proc      Production mechanism (VMC encoding)
     *@param ntr       Track number (filled by the stack)
     *@param weight    Particle weight
     *@param is        Generation status code (whatever that means)
     **/
    virtual void PushTrack(Int_t toBeDone, Int_t parentID, Int_t pdgCode,
                           Double_t px, Double_t py, Double_t pz,
                           Double_t e, Double_t vx, Double_t vy,
                           Double_t vz, Double_t time, Double_t polx,
                           Double_t poly, Double_t polz, TMCProcess proc,
                           Int_t& ntr, Double_t weight, Int_t is);

    virtual void PushTrack(Int_t toBeDone, Int_t parentID, Int_t pdgCode,
                           Double_t px, Double_t py, Double_t pz,
                           Double_t e, Double_t vx, Double_t vy,
                           Double_t vz, Double_t time, Double_t polx,
                           Double_t poly, Double_t polz, TMCProcess proc,
                           Int_t& ntr, Double_t weight, Int_t is, Int_t secondparentID);


    /** Virtual method PopNextTrack.
     ** Gets next particle for tracking from the stack.
     *@param  iTrack  index of popped track
     *@return Pointer to the TParticle of the track
     **/
    virtual TParticle* PopNextTrack(Int_t& iTrack);


    /** Virtual method PopPrimaryForTracking.
     ** Gets primary particle by index for tracking from stack.
     *@param  iPrim   index of primary particle
     *@return Pointer to the TParticle of the track
     **/
    virtual TParticle* PopPrimaryForTracking(Int_t iPrim);


    /** Add a TParticle to the fParticles array **/
    void AddParticle(TParticle* part);


    /** Fill the MCTrack output array, applying filter criteria **/
    virtual void FillTrackArray();


    /** Update the track index in the MCTracks and MCPoints **/
    virtual void UpdateTrackIndex(TRefArray* detArray);


    /** Resets arrays and stack and deletes particles and tracks **/
    virtual void Reset();


    /** Register the MCTrack array to the Root Manager  **/
    virtual void Register();


    /** Output to screen
     **@param iVerbose: 0=events summary, 1=track info
     **/
    virtual void Print(Int_t iVerbose=0) const;


    /** Modifiers  **/
    virtual void SetCurrentTrack(Int_t iTrack);

    /** Accessors **/
    virtual Int_t GetNtrack() const;   // Total number of tracks
    virtual Int_t GetNprimary() const; // Number of primaries
    virtual TParticle* GetCurrentTrack() const;
    virtual Int_t GetCurrentTrackNumber() const;
    virtual Int_t GetCurrentParentTrackNumber() const;
    virtual TParticle* GetParticle(Int_t trackID) const; 

    virtual void SelectTracks(Int_t, Int_t);

    /** Modifiers  **/
    void StoreSecondaries(Bool_t choice = kTRUE) { fStoreSecondaries = choice; }
    void SetMinPoints(Int_t min)                 { fMinPoints        = min;    }
    void SetEnergyCut(Double_t eMin)             { fEnergyCut        = eMin;   }
    void StoreMothers(Bool_t choice = kTRUE)     { fStoreMothers     = choice; }

    /** Accessors **/
    TClonesArray* GetListOfParticles() { return fParticles; }

  protected:

    template<class T>
      void FillTrackArray( Int_t, Int_t, T dataType );

    template<class T>
      void UpdateTrackIndex(TRefArray* detList, T dataType);

    template<class T>
      void Print(Int_t iVerbose, T dataType) const;

    /** STL stack (FILO) used to handle the TParticles for tracking **/
    std::stack<TParticle*>  fStack;           //!


    /** Array of TParticles (contains all TParticles put into or created
     ** by the transport
     **/
    TClonesArray* fParticles;            //!


    /** Array of FairMCTracks containg the tracks written to the output **/
    TClonesArray* fTracks;


    /** STL map from particle index to storage flag  **/
    std::map<Int_t, Bool_t>           fStoreMap;        //!
    std::map<Int_t, Bool_t>::iterator fStoreIter;       //!


    /** STL map from particle index to track index  **/
    std::map<Int_t, Int_t>            fIndexMap;        //!
    std::map<Int_t, Int_t>::iterator  fIndexIter;       //!


    /** STL map from track index and detector ID to number of MCPoints **/
    std::map<std::pair<Int_t, Int_t>, Int_t> fPointsMap;     //!


    /** Some indizes and counters **/
    Int_t fCurrentTrack;  //! Index of current track
    Int_t fNPrimaries;    //! Number of primary particles
    Int_t fNParticles;    //! Number of entries in fParticles
    Int_t fNTracks;       //! Number of entries in fTracks
    Int_t fIndex;         //! Used for merging


    /** Variables defining the criteria for output selection **/
    Bool_t     fStoreSecondaries;
    Int_t      fMinPoints;
    Double32_t fEnergyCut;
    Bool_t     fStoreMothers;



  private:
    FairGenericStack(const FairGenericStack&);
    FairGenericStack& operator=(const FairGenericStack&);

    ClassDef(FairGenericStack,1)


};

template<class T>
void FairGenericStack::FillTrackArray( Int_t begin, Int_t end, T datatype)
{

  LOG(INFO) << "Filling MCTrack array..." << FairLogger::endl;

  // --> Reset index map and number of output tracks
  fIndexMap.clear();
  fNTracks = 0;

  // --> Check tracks for selection criteria
  SelectTracks(begin, end);

  // --> Loop over fParticles array and copy selected tracks
  for (Int_t iPart=0; iPart<fNParticles; iPart++) {

    fStoreIter = fStoreMap.find(iPart);
    if (fStoreIter == fStoreMap.end() ) {
      LOG(FATAL) << "Particle " << iPart
		 << " not found in storage map!" << FairLogger::endl;
    }
    Bool_t store = (*fStoreIter).second;

    if (store) {
      T *track = (T*)fTracks->ConstructedAt(fNTracks);
      /*
      T* track =
        new( (*fTracks)[fNTracks]) T(GetParticle(iPart));
      */
      fIndexMap[iPart] = fNTracks;
      // --> Set the number of points in the detectors for this track
      for (Int_t iDet=begin; iDet<end; iDet++) {
	std::pair<Int_t, Int_t> a(iPart, iDet);
        track->SetNPoints(iDet, fPointsMap[a]);
      }
      fNTracks++;
    } else { fIndexMap[iPart] = -2; }

  }

  // --> Map index for primary mothers
  fIndexMap[-1] = -1;

  // --> Screen output
  Print(1);

}

template<class T>
void FairGenericStack::UpdateTrackIndex(TRefArray* detList, T dataType)
{

  LOG(INFO) << "Updating track indizes..." << FairLogger::endl;
  Int_t nColl = 0;

  // First update mother ID in MCTracks
  for (Int_t i=0; i<fNTracks; i++) {
    T* track = (T*)fTracks->At(i);
    Int_t iMotherOld = track->GetMotherId();
    fIndexIter = fIndexMap.find(iMotherOld);
    if (fIndexIter == fIndexMap.end()) {
      LOG(ERROR) << "Particle index " << iMotherOld
		 << " not found in dex map! " << FairLogger::endl;
    }
    track->SetMotherId( (*fIndexIter).second );
  }

  // Now iterate through all active detectors
  TIterator* detIter = detList->MakeIterator();
  detIter->Reset();
  FairDetector* det = NULL;
  while( (det = (FairDetector*)detIter->Next() ) ) {


    // --> Get hit collections from detector
    Int_t iColl = 0;
    TClonesArray* hitArray;
    while ( (hitArray = det->GetCollection(iColl++)) ) {
      nColl++;
      Int_t nPoints = hitArray->GetEntriesFast();

      // --> Update track index for all MCPoints in the collection
      for (Int_t iPoint=0; iPoint<nPoints; iPoint++) {
        FairMCPoint* point = (FairMCPoint*)hitArray->At(iPoint);
        Int_t iTrack = point->GetTrackID();

        fIndexIter = fIndexMap.find(iTrack);
        if (fIndexIter == fIndexMap.end()) {
          LOG(FATAL) << "Particle index " << iTrack
		     << " not found in index map! " << FairLogger::endl;
        }
        point->SetTrackID((*fIndexIter).second);
        point->SetLink(FairLink("MCTrack", (*fIndexIter).second));
      }

    }   // Collections of this detector
  }     // List of active detectors

  LOG(INFO) << "...stack and " << nColl << " collections updated." 
	    << FairLogger::endl;
}

template<class T>
void FairGenericStack::Print(Int_t iVerbose, T dataType) const
{
  LOG(INFO) << "Number of primaries        = "
	    << fNPrimaries << FairLogger::endl;
  LOG(INFO) << "Total number of particles  = "
	    << fNParticles << FairLogger::endl;
  LOG(INFO) << "Number of tracks in output = "
	    << fNTracks << FairLogger::endl;
  for (Int_t iTrack=0; iTrack<fNTracks; iTrack++) {
    ((T*) fTracks->At(iTrack))->Print(iTrack);
  }
}

#endif