/////////////////////////////////////////////////////////////////
//
// This class used to perform the GenFit for all track candidates.
//
/////////////////////////////////////////////////////////////////
#include <vector>
#include <iostream>

#include <IOADatabase.hxx>
#include <IGenFitter.hxx>
#include <IReconTrack.hxx>
#include <IReconTrackCand.hxx>

#include <ICOMETLog.hxx>
#include <ICOMETEvent.hxx>
#include <IHitSelection.hxx>
#include <IMCHit.hxx>
#include <IHit.hxx>
#include <IGeomInfo.hxx>
#include <IFieldManager.hxx>
#include <IGeoField.hxx>

#include <TCanvas.h>
#include <TH1F.h>
#include <TFile.h>
#include <TTree.h>
#include <TDirectory.h>

////// GenFit
#include <FieldManager.h>
#include <MaterialEffects.h>
#include <TGeoMaterialInterface.h>
#include <AbsMeasurement.h>
#include <WirePointMeasurement.h>
#include <PlanarMeasurement.h>
#include <SpacepointMeasurement.h>
#include <AbsTrackRep.h>
#include <RKTrackRep.h>
#include <StateOnPlane.h>
#include <MeasuredStateOnPlane.h>
#include <Track.h>
#include <KalmanFitter.h>
#include <KalmanFitterRefTrack.h>
#include <KalmanFitterInfo.h>
#include <KalmanFittedStateOnPlane.h>
#include <KalmanFitStatus.h>
#include <DAF.h>
#include <ConstField.h>

#include <TGeoManager.h>
#include <TGeoGlobalMagField.h>

/// For translating COMET Field into genfit::AbsBField
class GFGeoField : public genfit::AbsBField{
public:
  TVector3 get(const TVector3& pos) const{
    double B[3];
    get(pos.X(),pos.Y(),pos.Z(),B[0],B[1],B[2]);
    return TVector3(B[0],B[1],B[2]);
  }
  void get(const double& posX, const double& posY, const double& posZ, double& Bx, double& By, double& Bz) const{
    double xyz[3]={posX,posY,posZ},B[3];
    TGeoGlobalMagField::Instance()->Field(xyz,B);
    Bx=B[0];
    By=B[1];
    Bz=B[2];
    //std::cout<<posX<<" "<<posY<<" "<<posZ<<" "<<Bx<<" "<<By<<" "<<Bz<<std::endl;
  }
};

IGenFitter::IGenFitter(const char* name, const char* title)
 :IVTrackFitter(name, title)
 ,fUseBetheBloch(true)
 ,fUseBrems(false)
 ,fMethod("DAF")
 ,fPID(11)
 ,fMinIterations(10)
 ,fMaxIterations(50)
 ,fMinHitsInTrack(10)
 ,fMinNDF(4)
 ,fMaxMomentum(250)
 ,fMinMomentum(50)
 ,fMaxMomDiff(20)
 ,fUseExtGeomFile(false)
 ,fGeometry("")
 ,fUseExtFieldFile(false)
 ,fFieldMap("")
 ,fUseMCTruth(false)
 ,fSmearing(false)
 ,fSigmaD(0.02)
 ,fIsInitialized(false)
 ,fSaveHistogram(true)
 ,fSaveTree(false)
 ,fInitialMomentum(TVector3())
 ,fInitialPosition(TVector3())
 ,fFittedMomentum(TVector3())
 ,fFittedPosition(TVector3())
{
  if (fSaveTree) fTree = new TTree("gftree", "GenFit tree");
}

IGenFitter::~IGenFitter()
{
  if (fSaveTree) delete fTree;
}

int IGenFitter::Init()
{
  COMETInfo("//----------------------------------------------------------//");
  COMETInfo("// Initialize IGenFitter");

  if (fSaveHistogram) {
    fIniMomHist = new TH1F("IniMomHist", "Initial momentum", 200, 90, 110);
    fIniMomHist->SetXTitle("p_{true} [MeV]");
    fIniMomHist->SetYTitle("Entries");
    fFitMomHist = new TH1F("FitMomHist", "Fitted momentum", 200, 90, 110);
    fFitMomHist->SetXTitle("p_{fit} [MeV]");
    fFitMomHist->SetYTitle("Entries");
    fDifMomHist = new TH1F("DifMomHist", "p_{fit} - p_{true}", 500, -5, 5);
    fDifMomHist->SetXTitle("#Delta p [MeV]");
    fDifMomHist->SetYTitle("Entries");
    fDDistHist  = new TH1F("DDistHist",  "Drift distance", 240, 0, 1.2);
    fDDistHist->SetXTitle("drift distance [cm]");
    fDDistHist->SetYTitle("Entries");
    fNumHitHist = new TH1F("NumHitHist", "Number of hits", 125, 0, 125);
    fNumHitHist->SetXTitle("#of hits in track");
    fNumHitHist->SetYTitle("Entries");
    fChi2Hist   = new TH1F("Chi2Hist",   "#chi^{2}/ndf of fitting", 100, 0, 10);
    fChi2Hist->SetXTitle("#chi^{2}/ndf");
    fChi2Hist->SetYTitle("Entries");
    fMaxLayerHist = new TH1F("MaxLayerHist", "Maximum layers", 19, 0, 19);
    fMaxLayerHist->SetXTitle("# of layers");
    fMaxLayerHist->SetYTitle("Entries");
    fTOFLenHist = new TH1F("TOFLenHist", "TOF/TrkLen", 1000, 0., 0.25);
    fTOFLenHist->SetXTitle("(TOF)/(Track Length) [ns/cm]");
    fTOFLenHist->SetYTitle("Entries");
    fTOFLenHistSmear = new TH1F("TOFLenHistSmear", "TOF/TrkLen", 1000, 0., 0.25);
    fTOFLenHistSmear->SetXTitle("(TOF)/(Track Length) [ns/cm]");
    fTOFLenHistSmear->SetYTitle("Entries");
  }

  BeginOfEvent();

  COMETNamedInfo("IGenFitter", "//          Method = " << fMethod);
  COMETNamedInfo("IGenFitter", "// Fitting Parameters");
  COMETNamedInfo("IGenFitter", "//    MinIteration = " << fMinIterations);
  COMETNamedInfo("IGenFitter", "//    MaxIteration = " << fMaxIterations);
  COMETNamedInfo("IGenFitter", "//      BetheBloch = " << fUseBetheBloch);
  COMETNamedInfo("IGenFitter", "//           Brems = " << fUseBrems);
  COMETNamedInfo("IGenFitter", "// Use the Geometry in " << fGeometry.c_str());
  COMETNamedInfo("IGenFitter", "// Use the FiledMap in " << fFieldMap.c_str());
  COMETNamedInfo("IGenFitter", "//----------------------------------------------------------//");

  fIsInitialized = true;
  if (fSaveTree) {
    fTree->Branch("inimom", &fInitialMomentum);
    fTree->Branch("inipos", &fInitialPosition);
    fTree->Branch("fitmom", &fFittedMomentum);
    fTree->Branch("fitpos", &fFittedPosition);
  }
  return 0;
}

int IGenFitter::Finish() {
  COMETLog("Finish IGenFitter...");
  return 0;
}

int IGenFitter::BeginOfEvent()
{
  /// Field Map
  if (!genfit::FieldManager::getInstance()->isInitialized()) {
    if (fUseExtFieldFile) {
      /// Use COMET magnetic field
      COMET::IFieldManager::Import(fFieldMap.c_str());

      if (COMET::IFieldManager::getObject()) {
        TGeoGlobalMagField::Instance()->SetField(new COMET::IGeoField());
        COMETNamedInfo("IGenFitter", "IFieldManager found the field object");
      }
      else {
        TGeoGlobalMagField::Instance()->SetField(new TGeoUniformMagField(10,0,0)); /// 10 kilogauss corresponds to 1T and assuming that detector solenoid is along x-axis for Phase-I...
        COMETNamedInfo("IGenFitter", "constant field to be used");
      }

      /// Use TGeo field
      genfit::FieldManager::getInstance()->init(new GFGeoField());
    } else if (fIsInitialized) {
      COMET::IFieldManager::Import(); /// Open current field maps

      if (COMET::IFieldManager::getObject()) {
        TGeoGlobalMagField::Instance()->SetField(new COMET::IGeoField());
        COMETNamedInfo("IGenFitter", "IFieldManager found the field object");
      }
      else {
        TGeoGlobalMagField::Instance()->SetField(new TGeoUniformMagField(10,0,0));
        COMETNamedInfo("IGenFitter", "constant field to be used");
      }

      /// Use TGeo field
      genfit::FieldManager::getInstance()->init(new GFGeoField());
    }
  }

  /// Geometry & materials
  if (!genfit::MaterialEffects::getInstance()->isInitialized()) {
    if (fUseExtGeomFile) {
      /// Geometry
      TGeoManager::Import(fGeometry.c_str());

      /// Set material effects
      genfit::MaterialEffects* mateff = genfit::MaterialEffects::getInstance();
      mateff->setEnergyLossBetheBloch(fUseBetheBloch);
      mateff->setEnergyLossBrems(fUseBrems);

      genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
    } else if (fIsInitialized) {
      COMET::IOADatabase::Get().UpdateGeometry();

      /// Set material effects
      genfit::MaterialEffects* mateff = genfit::MaterialEffects::getInstance();
      mateff->setEnergyLossBetheBloch(fUseBetheBloch);
      mateff->setEnergyLossBrems(fUseBrems);

      genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
    }
  }
  fInitialMomentum = TVector3(1e9, 1e9, 1e9);
  fInitialPosition = TVector3(1e9, 1e9, 1e9);
  fFittedMomentum  = TVector3(1e9, 1e9, 1e9);
  fFittedPosition  = TVector3(1e9, 1e9, 1e9);

  return 0;
}

int IGenFitter::EndOfEvent()
{
  return 0;
}

COMET::IReconTrack* IGenFitter::doFit(const COMET::IHandle<COMET::IReconTrackCand>& aCand)
{
  genfit::Track *fitTrack(NULL);

  genfit::AbsTrackRep *rep = new genfit::RKTrackRep(fPID);

  /// Set initial position and state
  TVector3 posInit = aCand->GetSeedPosition();
  TVector3 momInit = aCand->GetSeedMomentum();

  /// Temporal solution to convert the unit
  posInit *= 0.1;   /// Convert from mm to cm
  momInit *= 0.001; /// Convert from MeV to GeV

  COMETNamedInfo("IGenFitter", "Initial position = (" <<  posInit.x() << ", " << posInit.y() << ", " << posInit.z() << ") cm" <<
                               ", momentum = (" << 1e3*momInit.x() << ", " << 1e3*momInit.y() << ". " << 1e3*momInit.z() << ") MeV" );

  fInitialMomentum  = momInit;
  fInitialMomentum *= 1e3;
  fInitialPosition  = posInit;

  genfit::MeasuredStateOnPlane stateInit(rep);
  TMatrixDSym covMInit(6);
  /// Position resolutions
  double resPos[3] = {0,0,0};
  resPos[0] = resPos[1] = resPos[2] = 0.1;  /// 1 mm
  for (int ii=0; ii<3; ii++) { covMInit(ii, ii) = resPos[ii]*resPos[ii]; }

  bool localCoord=false;

  //set covariant matrix from Momentum and angle resolution
  double sigma_p0     = 0.001; ///   1 MeV
  double sigma_angle0 = 0.1;   /// 100 mrad

  TVector3 dirP    = momInit.Unit();
  TVector3 zaxis   = localCoord ? TVector3(0,0,1) : TVector3(1,0,0);
  TVector3 v_axis  = dirP.Cross(zaxis).Unit();
  TVector3 u_axis  = dirP.Cross(v_axis).Unit();
  TVector3 axis[3] = {dirP, v_axis, u_axis};
  TMatrixD hrot(3, 3);
  for(int ii=0; ii<3; ii++) for(int jj=0; jj<3; jj++) hrot(ii,jj) = axis[jj](ii);
  TMatrixD cov0(3, 3);
  cov0(0,0) = pow(sigma_p0, 2.);
  cov0(1,1) = pow(sigma_angle0*momInit.Mag(), 2.);
  cov0(2,2) = pow(sigma_angle0*momInit.Mag(), 2.);
  TMatrixD covrot(3, 3);
  covrot    = hrot*cov0;
  covrot   *= hrot.T();

  for(int ii=0; ii<3; ii++) for(int jj=0; jj<3; jj++) covMInit[ii+3][jj+3] = covrot(ii, jj);

  rep->setPosMomCov(stateInit, posInit, momInit, covMInit);

  /// Set Track
  TVectorD seedState(6);
  TMatrixDSym seedCov(6);
  rep->get6DStateCov(stateInit, seedState, seedCov);
  fitTrack = new genfit::Track(rep, seedState, seedCov);

  std::vector<genfit::TrackPoint*>     trackPoints; // each track point
  std::vector<genfit::AbsMeasurement*> mesHits;     // each measured hit

  /// Fill hits
  const COMET::IHitSelection* hits = &(*(aCand->GetHits()));
  if (hits->size()<fMinHitsInTrack) {
    COMETNamedInfo("IGenFitter", "Number of hits are too small :" << hits->size());
    delete fitTrack;
    return NULL;
  }

  unsigned int nHits =  0; // number of total hits
  int numMaxLayer    = -1;
  double prevZ       =  0.;
  int prevStation    = -1;
  int nVirtualPlanes =  0;
  int nTotalHits     =  0;
  double t0          =  -1e9;
  double tECAL       =  -1e9;
  double tTOF        =  -1e9;
  COMET::IChannelId  chanId;
  COMET::IGeometryId geomId;

  bool isStrawTrack = false;
  bool isCDCTrack   = false;

  for (COMET::IHitSelection::const_iterator hitIt = hits->begin(); hitIt != hits->end(); ++hitIt) {

    chanId = (*hitIt)->GetChannelId();
    geomId = (*hitIt)->GetGeomId();
    if (t0<0) t0 = (*hitIt)->GetTime();

    if (fUseMCTruth) {
      UInt_t channel = chanId.AsUInt();
      trackPoints.push_back(new genfit::TrackPoint());
      TVectorD pos(3);
      TMatrixDSym posMatrix(3);
      COMET::IHandle<COMET::IMCHit> mcHit = (*hitIt);
      pos[0] = mcHit->GetPositionMC().X();
      pos[1] = mcHit->GetPositionMC().Y();
      pos[2] = mcHit->GetPositionMC().Z();
      if (fSmearing) { /// smear with same weight
        for (int ii=0;ii<3;ii++) pos[ii] += gRandom->Gaus(0, fSigmaD/TMath::Sqrt(3.));
      }

      for (int row = 0; row < 3; row++) {
        for (int col = 0; col < 3; col++) {
          if (row==col) posMatrix(row,col) = std::pow(fSigmaD/TMath::Sqrt(3.),2);
          else posMatrix(row,col) = 0;
        }
      }

      mesHits.push_back( new genfit::SpacepointMeasurement(pos, posMatrix, channel, nHits, trackPoints.at(nHits)) );
    }
    else if (chanId.IsECALChannel()) { /// Add ECAL hit
      UInt_t channel = chanId.AsUInt();
      if (tECAL<0 || tECAL>(*hitIt)->GetTime()) { // use only the 1st conversion point...
        trackPoints.push_back(new genfit::TrackPoint());
        tECAL = (*hitIt)->GetTime();
        TVectorD pos(3);
        TMatrixDSym posMatrix(3);
        COMET::IHandle<COMET::IMCHit> mcHit = (*hitIt);
        pos[0] = mcHit->GetPositionMC().X();
        pos[1] = mcHit->GetPositionMC().Y();
        pos[2] = mcHit->GetPositionMC().Z();
        posMatrix(0,0) = std::pow(2.5,2.);
        posMatrix(1,1) = std::pow(1.5,2.);
        posMatrix(2,2) = std::pow(1.5,2.);
        for (int row = 0; row < 3; row++) {
          for (int col = 0; col < 3; col++) {
            if (row!=col) posMatrix(row,col) = 0.;
          }
        }
        mesHits.push_back( new genfit::SpacepointMeasurement(pos, posMatrix, channel, nHits, trackPoints.at(nHits)) );
      } else continue;
    }
    else if (chanId.IsTOFChannel()) { /// Add TOF hit
      UInt_t channel = chanId.AsUInt();
      if (tTOF<0) { // use only the 1st conversion point...
        trackPoints.push_back(new genfit::TrackPoint());
        tTOF = (*hitIt)->GetTime();
        TVectorD pos(3);
        TMatrixDSym posMatrix(3);
        COMET::IHandle<COMET::IMCHit> mcHit = (*hitIt);
        pos[0] = mcHit->GetPositionMC().X();
        pos[1] = mcHit->GetPositionMC().Y();
        pos[2] = mcHit->GetPositionMC().Z();
        posMatrix(0,0) = std::pow(0.01,2.);
        posMatrix(1,1) = std::pow(0.25,2.);
        posMatrix(2,2) = std::pow(0.25,2.);
        for (int row = 0; row < 3; row++) {
          for (int col = 0; col < 3; col++) {
            if (row!=col) posMatrix(row,col) = 0.;
          }
        }
        mesHits.push_back( new genfit::SpacepointMeasurement(pos, posMatrix, channel, nHits, trackPoints.at(nHits)) );
      } else continue;
    }
    else if (chanId.IsCDCChannel()) { /// Add CDC hit
      isCDCTrack = true;

      trackPoints.push_back(new genfit::TrackPoint());

      int wire = COMET::IGeomInfo::Get().CDC().GeomIdToWire(geomId);
      TVectorD wireMes(7);       /// Wire end0(x,y,z), end1(x,y,z), drift distance
      TMatrixDSym wireMatrix(7); /// Uncertainties for wireMes values

      wireMes[0] = COMET::IGeomInfo::Get().CDC().GetWireEnd0(wire).X();
      wireMes[1] = COMET::IGeomInfo::Get().CDC().GetWireEnd0(wire).Y();
      wireMes[2] = COMET::IGeomInfo::Get().CDC().GetWireEnd0(wire).Z();
      wireMes[3] = COMET::IGeomInfo::Get().CDC().GetWireEnd1(wire).X();
      wireMes[4] = COMET::IGeomInfo::Get().CDC().GetWireEnd1(wire).Y();
      wireMes[5] = COMET::IGeomInfo::Get().CDC().GetWireEnd1(wire).Z();
      wireMes[6] = (*hitIt)->GetDriftDistance();
      if (fSmearing) wireMes[6] += gRandom->Gaus(0, fSigmaD); /// drift distance is smeared by fSigmaD

      if (numMaxLayer<COMET::IGeomInfo::Get().CDC().GetLayer(wire))
        numMaxLayer = COMET::IGeomInfo::Get().CDC().GetLayer(wire);

      if (fSaveHistogram) fDDistHist->Fill(wireMes[6]);

      COMETNamedDebug("IGenFitter", " wire = " << wire << ", wire end0 = " << wireMes[0] << " " << wireMes[1] << " " << wireMes[2]
                     << ", wire end1 = " << wireMes[3] << " " << wireMes[4] << " " << wireMes[5] << ", drift distance = " << wireMes[6]);

      for (int row = 0; row < 7; row++) {
        for (int col = 0; col < 7; col++) {
          if (row != col) wireMatrix(row,col) = 0;
          if (row < 6)    wireMatrix(row,col) = std::pow(0.001,2);   /// Wire position uncertainties (temporary 10um)
          else            wireMatrix(row,col) = std::pow(fSigmaD,2); /// Resolution of drift distance
        }
      }
      genfit::WireMeasurement *whit;
      mesHits.push_back(whit=new genfit::WireMeasurement(wireMes, wireMatrix, wire, nHits, trackPoints.at(nHits)) );
      whit->setLeftRightResolution((*hitIt)->GetLeftRight());
    }
    else if (chanId.IsStrawTrkChannel()) { /// Add StrawTrk hit
      isStrawTrack = true;

      trackPoints.push_back(new genfit::TrackPoint());
      TVectorD wireMes(7);       /// Wire end0(x,y,z), end1(x,y,z), drift distance
      TMatrixDSym wireMatrix(7); /// Uncertainties for wireMes values

      COMET::IGeomInfo& geomInfo = COMET::IGeomInfo::Get();
      Int_t station   = geomInfo.StrawTrk().GetStationIdFromGeomId(geomId);
      Int_t manifoldId = geomInfo.StrawTrk().GetManifoldIdFromGeomId(geomId);
      Int_t straw     = geomInfo.StrawTrk().GetStrawIdFromGeomId(geomId);
      Int_t wire      = geomInfo.StrawTrk().GetSequenceIdFromStrawTrk(station,manifoldId,straw);

      wireMes[0] = geomInfo.StrawTrk().GetWireEnd0(station, manifoldId, straw).X();
      wireMes[1] = geomInfo.StrawTrk().GetWireEnd0(station, manifoldId, straw).Y();
      wireMes[2] = geomInfo.StrawTrk().GetWireEnd0(station, manifoldId, straw).Z();
      wireMes[3] = geomInfo.StrawTrk().GetWireEnd1(station, manifoldId, straw).X();
      wireMes[4] = geomInfo.StrawTrk().GetWireEnd1(station, manifoldId, straw).Y();
      wireMes[5] = geomInfo.StrawTrk().GetWireEnd1(station, manifoldId, straw).Z();
      wireMes[6] = (*hitIt)->GetDriftDistance();
      if (fSmearing) wireMes[6] += gRandom->Gaus(0, fSigmaD); /// drift distance is smeared by fSigmaD

      COMETNamedDebug("IGenFitter", "stationId, manifoldId, wireId = (" << station << ", " << manifoldId <<", "<< straw << "), " <<
                     "wire end0 = (" << wireMes[0] << ", " << wireMes[1] << ", " << wireMes[2] << "), " <<
                     "wire end1 = (" << wireMes[3] << ", " << wireMes[4] << ", " << wireMes[5] << "), drift distance = " << wireMes[6]);

      if(prevStation!=-1 && station!=prevStation){
        TVectorD hitCoords(2);
        TMatrixDSym hitCov(2);
        hitCov(0,0) = 1e30;
        hitCov(1,1) = 1e30;

        TVector3 wireAxis = (zaxis.Cross(TVector3(0,1,0))).Cross(zaxis);
        genfit::PlanarMeasurement* measurement
            = new genfit::PlanarMeasurement(hitCoords, hitCov, -1, nTotalHits, nullptr);
        genfit::SharedPlanePtr plane(new genfit::DetPlane(zaxis*(prevZ+1.),
                                                          zaxis.Cross(TVector3(0,1,0)),
                                                          wireAxis));
        measurement->setPlane(plane);
        nVirtualPlanes++;
        fitTrack->insertMeasurement(measurement, nTotalHits++);

        measurement
            = new genfit::PlanarMeasurement(hitCoords, hitCov, -1, nTotalHits, nullptr);
        genfit::SharedPlanePtr plane2(new genfit::DetPlane(zaxis*((localCoord?wireMes[2]:wireMes[0])-1.),
                                                           zaxis.Cross(TVector3(0,1,0)),
                                                           wireAxis));
        measurement->setPlane(plane2);
        nVirtualPlanes++;
        fitTrack->insertMeasurement(measurement, nTotalHits++);
      }
      prevStation = station;
      prevZ = (localCoord ? wireMes[2] : wireMes[0]);

      if (fSaveHistogram) fDDistHist->Fill(wireMes[6]);

      for (int row = 0; row < 7; row++) {
        for (int col = 0; col < 7; col++) {
          if (row != col) wireMatrix(row,col) = 0;
          if (row < 6)    wireMatrix(row,col) = std::pow(0.001, 2);    /// Wire position uncertainties (temporary 10um)
          else            wireMatrix(row,col) = std::pow(fSigmaD, 2);  /// Resolution of drift distance
        }
      }
      genfit::WireMeasurement *whit;
      mesHits.push_back(whit=new genfit::WireMeasurement(wireMes, wireMatrix, wire, nTotalHits, trackPoints.at(nHits)) );
      whit->setLeftRightResolution((*hitIt)->GetLeftRight());
    }
    else {
      COMETNamedInfo("IGenFitter", "Not CDC or StrawTracker hit");
      continue;
    }
/*
    if (hitType == kPlanarMeasurement) { /// PlanarMeasurement to be used by CTH?
      TVectorD planeMes;
      TMatrixDSym planeMatrix;

      mesHits.at(nHits) = new genfit::PlanarMeasurement(planeMes, planeMatrix, detId, nHits, trackPoints.at(nHits));
    } else if (hitType == kSpacepointMeasurement) { /// SpacepointMeasurement to be used by ECAL?
      TVectorD pointMes;
      TMatrixDSym pointMatrix;

      mesHits.at(nHits) = new genfit::SpacepointMeasurement(pointMes, pointMatrix, detId, nHits, trackPoints.at(nHits));
    } else {

      COMETLog("This hit type is invalid: " << hitType);
      continue;
    }
*/
    trackPoints.at(nHits)->addRawMeasurement(mesHits.at(nHits));
    fitTrack->insertPoint(trackPoints.at(nHits), nTotalHits++);
    nHits = mesHits.size();
  }
  numMaxLayer += 1;

  if (nHits<fMinHitsInTrack) {
    COMETNamedInfo("IGenFitter", "Number of hits are too small :" << nHits);
    delete fitTrack;
    return NULL;
  }

  /// Set a fitter
  genfit::AbsKalmanFitter *kalman = NULL;
  if (fMethod=="DAF") {
    kalman = new genfit::DAF();
  } else if (fMethod=="KalmanFitter") {
    kalman = new genfit::KalmanFitter(fMaxIterations);
    kalman->setMinIterations(fMinIterations);
  } else if (fMethod=="KalmanFitterRefTrack") {
    kalman = new genfit::KalmanFitterRefTrack(fMaxIterations);
    kalman->setMinIterations(fMinIterations);
  } else {
    COMETNamedInfo("IGenFitter", "This method is invalid: " << fMethod);
    delete fitTrack;
    return NULL;
  }

  /// Do the fitting
  try{
    kalman->processTrack(fitTrack);
  }catch(genfit::Exception& e){e.what();}

  if (!fitTrack->getFitStatus(rep)->isFitConverged()) {
    kalman->processTrack(fitTrack);
  }

  if (!fitTrack->getFitStatus(rep)->isFitted()||!fitTrack->getFitStatus(rep)->isFitConverged()) {

    COMETNamedInfo("IGenFitter", "Fitting is failed...");

    delete kalman;
    delete fitTrack;
    return NULL;
  }

  if (fitTrack->getFitStatus(rep)->getChi2()<=0 ||
      (fitTrack->getFitStatus(rep)->getNdf()-2*nVirtualPlanes)<fMinNDF) {
    COMETNamedInfo("IGenFitter", "Fit result might be wrong... (chi2,ndf) = (" << 
                   fitTrack->getFitStatus(rep)->getChi2() << "," << 
                   (fitTrack->getFitStatus(rep)->getNdf()-2*nVirtualPlanes) << ")");
    delete kalman;
    delete fitTrack;
    return NULL;
  }

  if (fSaveHistogram) {
    fNumHitHist->Fill(nTotalHits);
    fMaxLayerHist->Fill(numMaxLayer);
  }
  COMET::IReconTrack* reconTrack
    = CopyATrack(fitTrack, rep, nTotalHits, nVirtualPlanes, isStrawTrack);
  if (isStrawTrack && tECAL>0. && tTOF) {
    try {
      Double_t trkLen = fitTrack->getTrackLen(rep, 2, nTotalHits-1);
      COMETNamedInfo("IGenFitter", "TOF/TrkLen = " << tECAL-tTOF << "/" << trkLen << " = " << (tECAL-tTOF)/trkLen);
      fTOFLenHist->Fill((tECAL-tTOF)/trkLen);
      tECAL += gRandom->Gaus(0, 0.6); // 600ps smearing
      tTOF  += gRandom->Gaus(0, 1.0); //   1ns smearing
      COMETNamedInfo("IGenFitter", "TOF/TrkLen = " << tECAL-tTOF << "/" << trkLen << " = " << (tECAL-tTOF)/trkLen << " AFTER Smearing");
      fTOFLenHistSmear->Fill((tECAL-tTOF)/trkLen);
    }catch(genfit::Exception& e){e.what();}
  }

  delete kalman;
  delete fitTrack;

  COMETNamedDebug("IGenFitter", " track momentum = " << reconTrack->GetMomentum());
  return reconTrack;
}


COMET::IHandle<COMET::IAlgorithmResult> IGenFitter::Process(const COMET::IAlgorithmResult& input)
{
  COMETNamedInfo("IGenFitter", "IGenFitter::Process");
  COMET::IOADatabase::Get().UpdateGeometry();
  /// Output algorithm result
  COMET::IHandle<COMET::IAlgorithmResult>
    result(new COMET::IAlgorithmResult(Form("IGenFitResult_%s",fName.Data()), Form("GenFitTracks for %s",fName.Data())));


  COMET::IHandle<COMET::IReconObjectContainer>
    inputTrks(input.GetResultsContainer("mctrackcandidates")); // to be modified in order to read every possible track candidate containers

  if(!inputTrks) {
    COMETNamedInfo("IGenFitter", "No inputs");
    return result;
  }

  /// Make track container
  COMET::IReconObjectContainer*
    tracks = new COMET::IReconObjectContainer(Form("GenFit_TrackContainer_%s",fName.Data()), Form("Container of GenFit Tracks for %s",fName.Data()));

  /// A candidate track
  COMET::IHandle<COMET::IReconTrackCand> aCand;

  COMET::IReconTrack* aTrack;

  if (!inputTrks->size()) COMETNamedInfo("IGenFitter", "No input ReconObject");

  for (COMET::IReconObjectContainer::iterator trk = inputTrks->begin(); trk != inputTrks->end(); ++trk) {
    aTrack = NULL;

    aCand = (*trk);

    if (!aCand) continue;

    aTrack = doFit(aCand);

    if (aTrack) tracks->push_back(COMET::IHandle<COMET::IReconTrack>(aTrack));
  }

  /// Add to the result
  if (tracks->size() > 0){
    result->AddResultsContainer(tracks);
    COMETNamedDebug("IGenFitter", "AlgorithmResult name = " << result->GetName() << " track size = " << tracks->size());
  }
  else {
    COMETNamedInfo("IGenFitter", "No tracks");
    delete tracks;
  }

  return result;
}

/// Genfit track represantation will be translated to the ReconTrack object here.
COMET::IReconTrack* IGenFitter::CopyATrack(const genfit::Track* gtrack, genfit::AbsTrackRep* rep,
                                           const int nHits, const int nVirtualHits,
                                           const bool ecalExtrapolate)
{
  double chi2 = gtrack->getFitStatus(rep)->getChi2();
  double ndf  = gtrack->getFitStatus(rep)->getNdf()/*-2*nVirtualHits*/;
  COMETNamedDebug("IGenFitter", "CopyATrack-->" << " chi2/ndf = " << chi2/ndf);

  //genfit::MeasurementOnPlane   res; // not used now
  const std::vector<genfit::TrackPoint*> points = gtrack->getPoints();
  COMET::IReconTrack* reconTrack = new COMET::IReconTrack();
  Int_t posStateIdx, dirStateIdx, momStateIdx;
  posStateIdx = dirStateIdx = momStateIdx = -1;
  Double_t tmpT = 0., tmpL = 0.;
  std::vector<Double_t> time;
  std::vector<Double_t> length;
  std::vector<Int_t> usedHit;
  TVector3    posOnPlane;
  TVector3    momOnPlane;
  TMatrixDSym covOnPlane(6);
  for (int iHit = 0; iHit < nHits; iHit++) {
    // get fitted status of fit track
    int hitId = (points.at(iHit))->getRawMeasurement(0)->getHitId();
    int detId = (points.at(iHit))->getRawMeasurement(0)->getDetId();
    if(hitId<0) continue;

    genfit::MeasuredStateOnPlane mop;
    if (points.at(iHit) != NULL) {
      mop = gtrack->getFittedState(iHit,rep);
      /*if ((points.at(iHit))->hasFitterInfo(rep)) {
        res = (points.at(iHit))->getKalmanFitterInfo(rep)->getResidual(); // not used now
      }*/
      try{
        mop.getPosMomCov(posOnPlane,momOnPlane,covOnPlane);
      } catch(genfit::Exception& e){ e.what();continue; }
    } else continue;

    if (iHit==0) {
      fFittedMomentum  = momOnPlane;
      fFittedMomentum *= 1e3;
      fFittedPosition  = posOnPlane;
      COMETNamedDebug("IGenFitter", " Chi2/NDF          = "
                      << chi2/ndf << " ( " << chi2 << " / " << ndf << " ) " << " nHits " << nHits);
      COMETNamedDebug("IGenFitter", " Fitted momentum   = " << fFittedMomentum.Mag() << " MeV");
      COMETNamedDebug("IGenFitter", " Momentum Residual = " << fFittedMomentum.Mag()-fInitialMomentum.Mag() << " MeV");
      if (fSaveHistogram) {
        fIniMomHist->Fill(fInitialMomentum.Mag());
        fFitMomHist->Fill(fFittedMomentum.Mag());
        fDifMomHist->Fill(fFittedMomentum.Mag()-fInitialMomentum.Mag());
        fChi2Hist->Fill(chi2/ndf);
      }
      if (fSaveTree) fTree->Fill();

      time.push_back(0);
      length.push_back(0);
      COMET::IHandle<COMET::ITrackState> trkState = reconTrack->GetState();

      if (posStateIdx<0) posStateIdx = trkState->GetPositionIndex();
      if (dirStateIdx<0) dirStateIdx = trkState->GetDirectionIndex();
      if (momStateIdx<0) momStateIdx = trkState->GetMomentumIndex();

      SetTrackState(trkState, momOnPlane, posOnPlane, time.back());
      SetTrackCovariance(trkState, covOnPlane, posStateIdx, dirStateIdx, momStateIdx);

    } else {
      if (iHit!=(nHits-1)) { /// getTOF abort when iHit==nHits-1 ...
        try{
          tmpT = gtrack->getTOF(rep, 0, iHit);
          tmpL = gtrack->getTrackLen(rep, 0, iHit);
        } catch (genfit::Exception& e) { e.what(); break; }
        time.push_back(tmpT);
        length.push_back(tmpL);
      } else { /// Update TOF using previous 2 points
        try{
          tmpT = time.back() + (gtrack->getTrackLen(rep, 0, iHit)-length.back())
                               *(time.back()-time.at(time.size()-3))/(length.back()-length.at(length.size()-3));
          tmpL = gtrack->getTrackLen(rep, 0, iHit);
        } catch (genfit::Exception& e) { e.what(); break; }
        time.push_back(tmpT);
        length.push_back(tmpL);
      }
      COMETNamedDebug("IGenFitter", "Mom @ Hit " << iHit << ": "
                      << momOnPlane.X() << " " << momOnPlane.Y() << " " << momOnPlane.Z());
      COMETNamedDebug("IGenFitter", "Pos @ Hit " << iHit <<  ": "
                      << posOnPlane.X() << " " << posOnPlane.Y() << " " << posOnPlane.Z());

      COMET::IReconTrack* trkPoint = new COMET::IReconTrack();
      COMET::IHandle<COMET::ITrackState> trkState = trkPoint->GetState();
      SetTrackState(trkState, momOnPlane, posOnPlane, time.back());
      SetTrackCovariance(trkState, covOnPlane, posStateIdx, dirStateIdx, momStateIdx);

      COMET::IHandle<COMET::IReconNode> aNode(new COMET::IReconNode());
      COMET::IHandle<COMET::IReconBase> aBase(trkPoint);
      aNode->SetObject(aBase);
      reconTrack->GetNodes().push_back(aNode);
    }
    usedHit.push_back(iHit);
  }
/*
  if (ecalExtrapolate && usedHit.size()>0) {
    //genfit::TrackPoint *aPoint = gtrack->getPointWithMeasurementAndFitterInfo(nHits-1, rep);
    genfit::TrackPoint *aPoint = gtrack->getPointWithMeasurementAndFitterInfo(usedHit.back(), rep);
    if (aPoint!=NULL) { /// Extrapolate to the ECAL surface plane
      TVector3 ecalOrig = TVector3(839.5, 0, 765.); /// Only for this moment, this has to be the original position of ECAL surface
      genfit::KalmanFittedStateOnPlane stateOnECAL(*(static_cast<genfit::KalmanFitterInfo*>(aPoint->getFitterInfo(rep))->getForwardUpdate()));
      // ECLA surface plane
      genfit::SharedPlanePtr ecalPlane(new genfit::DetPlane(ecalOrig,
                                                            TVector3(0, 1, 0),
                                                            TVector3(0, 0,-1)));
      try{
        rep->extrapolateToPlane(stateOnECAL, ecalPlane);
      }catch(genfit::Exception& e){e.what();}

      TVectorD&  stateECAL = stateOnECAL.getState();
      TMatrixDSym& covECAL = stateOnECAL.getCov();
      genfit::MeasuredStateOnPlane* mop = new genfit::MeasuredStateOnPlane(stateECAL, covECAL, ecalPlane, rep);

      mop->getPosMomCov(posOnPlane, momOnPlane, covOnPlane);
      time.push_back(1e9); // only for this moment
      COMETNamedDebug("IGenFitter", "Mom @ ECAL : "
                      << momOnPlane.X() << " " << momOnPlane.Y() << " " << momOnPlane.Z());
      COMETNamedDebug("IGenFitter", "Pos @ ECAL : "
                      << posOnPlane.X() << " " << posOnPlane.Y() << " " << posOnPlane.Z());

      /// Save the track state at ECAL surface plane...
      COMET::IReconTrack* trkPoint = new COMET::IReconTrack();
      COMET::IHandle<COMET::ITrackState> trkState = trkPoint->GetState();
      SetTrackState(trkState, momOnPlane, posOnPlane, time.back());

      COMET::IHandle<COMET::IReconNode> aNode(new COMET::IReconNode());
      COMET::IHandle<COMET::IReconBase> aBase(trkPoint);
      aNode->SetObject(aBase);
      reconTrack->GetNodes().push_back(aNode);

      delete mop;
    }
  }
*/
  return reconTrack;
}

void IGenFitter::SetTrackState(COMET::IHandle<COMET::ITrackState>& trackState,
                               const TVector3& mom, const TVector3& pos, const double& time)
{
  trackState->SetMomentum(mom.Mag());
  trackState->SetDirection(mom); /// Don't normalize to keep consistency with the values in covariance matrix
  trackState->SetPosition(TLorentzVector(pos.X(), pos.Y(), pos.Z(), time));
}

void IGenFitter::SetTrackCovariance(COMET::IHandle<COMET::ITrackState>& trackState, const TMatrixDSym& cov,
                                    const int posIdx, const int dirIdx, const int momIdx)
{
  for (Int_t i = 0; i < 3; ++i)
    for (Int_t j = 0; j < 3; ++j)
      trackState->SetCovarianceValue(i+posIdx, j+posIdx, cov(i,j));

  for (Int_t i = 0; i < 3; ++i)
    for (Int_t j = 0; j < 3; ++j)
      trackState->SetCovarianceValue(i+dirIdx, j+dirIdx, cov(i+3,j+3));

  trackState->SetCovarianceValue(momIdx, momIdx, cov(3,3)+cov(4,4)+cov(5,5));
}

int IGenFitter::Load(COMET::ICOMETEvent& anEvent)
{
  return 1;
}

int IGenFitter::Save(COMET::ICOMETEvent* anEvent)
{
  return 1;
}