/////////////////////////////////////////////////////////////////
//
// This class used to perform the GenFit for all track candidates.
// written by Y.Fujii (yfujii@ihep.ac.cn)
//
/////////////////////////////////////////////////////////////////


#ifndef TGenFitter_hxx_seen
#define TGenFitter_hxx_seen

#include <vector>

#include <TObject.h>
#include <TVector3.h>
#include <TH1F.h>

#include <IAlgorithm.hxx>
#include <ICOMETLog.hxx>

#include <IVTrackFitter.hxx>

#include <IHandle.hxx>
#include <ICOMETEvent.hxx>
#include <IReconTrack.hxx>
#include <IReconTrackCand.hxx>
#include <IFieldManager.hxx>

#include <TVirtualFitter.h>
#include <TRandom.h>

////// GenFit
#include <Track.h>

class IGenFitter : public IVTrackFitter {
public:
    IGenFitter(const char* name, const char* title);
    virtual ~IGenFitter();

    /// Main fitting function called for each track candidate
    COMET::IReconTrack* doFit(const COMET::IHandle<COMET::IReconTrackCand>& aCand);

    /// Main task called by ICOMETReconEventFunction
    COMET::IHandle<COMET::IAlgorithmResult> Process(const COMET::IAlgorithmResult& input);

    /// load the event
    virtual Int_t  Load(COMET::ICOMETEvent& anEvent);

    /// save the event
    virtual Int_t  Save(COMET::ICOMETEvent* anEvent);

    /// called at the begin of run or else (should not be in event-by-event)
    virtual Int_t  Init();

    /// called at the end of run or else (should not be in event-by-event)
    virtual Int_t  Finish();

    /// called at the begin of each event
    virtual Int_t  BeginOfEvent();

    /// called at the end of each event
    virtual Int_t  EndOfEvent();

    /// copy a track
    COMET::IReconTrack* CopyATrack(const genfit::Track* gtrack, genfit::AbsTrackRep* rep,
                                   const Int_t nHits, const Int_t nVirtualHits,
                                   const Bool_t ecalExtrapolate);

    /// Set the name of geometry file
    void SetGeometry(std::string name) {fGeometry = name;}

    /// Set the name of filedmap file
    void SetFieldMap(std::string name) {fFieldMap = name;}

    /// Set flag for using MC truth hit position
    void SetUseMCTruth(Bool_t val) {fUseMCTruth = val;}

    /// Set flag for the drift distance smearing
    void SetSmearing(Bool_t val) {fSmearing = val;}

    /// Set the position(drift distance) resolution
    void SetSigmaDistance(Double_t val) {fSigmaD = val;}

    /// Set the Track State
    void SetTrackState(COMET::IHandle<COMET::ITrackState>& trackState,
                       const TVector3& mom, const TVector3& pos, const Double_t& time);

    /// Set the Track Covariance
    void SetTrackCovariance(COMET::IHandle<COMET::ITrackState>& trackState, const TMatrixDSym& cov,
                            const Int_t posIdx, const Int_t dirIdx, const Int_t momIdx);

    void SetPID(Int_t val) {fPID = val;}

    Int_t GetPID() {return fPID;}

    enum {
      kWireMeasurement,
      kPlanarMeasurement,
      kSpacepointMeasurement
    };

protected:

    TH1F*  fIniMomHist;   /// Initial momentum

    TH1F*  fFitMomHist;   /// Fitted momentum

    TH1F*  fDifMomHist;   /// Fitted - Initial momentum

    TH1F*  fNumHitHist;   /// Number of hits in track

    TH1F*  fChi2Hist;     /// Chi square of fitting

    TH1F*  fDDistHist;    /// Drift distance

    TH1F*  fMaxLayerHist; /// Maximum layers

    TH1F*  fTOFLenHist;   /// TOF/TrkLen [ns/cm]

    TH1F*  fTOFLenHistSmear;   /// TOF/TrkLen [ns/cm]

    Bool_t fUseBetheBloch; /// flag to turn on/off the BetheBloch

    Bool_t fUseBrems;      /// flag to turn on/off the Brems

    std::string fMethod; /// fitting method

    Int_t  fPID;           /// Input Particle ID in PDGEncoding to fit the track

    UInt_t fMinIterations; /// minimum number of iterations

    UInt_t fMaxIterations; /// maximum number of iterations

    UInt_t fMinHitsInTrack;   /// minimum number of hits in track

    Int_t  fMinNDF;           /// minimum value for NDF

    Double_t fMaxMomentum;   /// maximum momentum [MeV]

    Double_t fMinMomentum;   /// minimum momentum [MeV]

    Double_t fMaxMomDiff;    /// maximum allowed difference between fitted and initial momenta [MeV]

    Bool_t   fUseExtGeomFile;/// flag to use geometry in external file

    std::string fGeometry; /// name of the geometry

    Bool_t   fUseExtFieldFile; /// flag to use field maps in external file

    std::string fFieldMap; /// name of the fieldmap

    Bool_t   fUseMCTruth;    /// Flag to use MC true hit position

    Bool_t   fSmearing;      /// enable Gaussian smearing for the drift distance using given fSigmaD

    Double_t fSigmaD;        /// position resolution

    Bool_t   fIsInitialized; /// flag to check this class is initialized or not.

    Bool_t   fSaveHistogram; /// flag to save the histograms

    Bool_t   fSaveTree;      /// flag to save tree

    TTree   *fTree;
    TVector3 fInitialMomentum;
    TVector3 fInitialPosition;
    TVector3 fFittedMomentum;
    TVector3 fFittedPosition;
};
#endif