#define __DEBUG__

#include "TrackFitterKalman.h"
#include "TrackFitterFactory.h"
#include <marlin/StringParameters.h>
#include <EVENT/LCRunHeader.h>

#include <streamlog/streamlog.h>

#include "LCObjectCopier.h"

// GEAR
#include "gear/GEAR.h"
#include "gear/TPCParameters.h"
#include "gear/PadRowLayout2D.h"
#include "gearimpl/TPCModuleImpl.h"
#include "gearxml/GearXML.h"

// global constants from Marlin, used for the global pointer to the GearMgr
#include <marlin/Global.h>

// KalTest
#include "TKalTrackSite.h"
#include "TKalTrackState.h"
#include "TKalDetCradle.h"
#include "THelicalTrack.h"

// KalDetector, MeasLayer and KalHit
#include "EXTPCKalDetector.h"
#include "EXTPCMeasLayer.h"
#include "EXTPCHit.h"
#include "EXHYBTrack.h"

// ROOT
#include "TVector3.h"
#include "TMath.h"

#if 1
#include "TFile.h"
#include "TNtupleD.h"
#include "TH2.h"
#include "TH1.h"
#endif


// Others
#include <sstream>
#include <iomanip>
#include <complex>

using namespace std;

namespace marlintpc {
  static const Bool_t gkDir = kIterBackward;

#if 1
  // use in estimation of tracking efficiency
  Bool_t TrackFitterKalman::fgIsUpperHit = kFALSE;
  Bool_t TrackFitterKalman::fgIsLowerHit = kFALSE;
  Bool_t TrackFitterKalman::fgIsMidHit   = kFALSE;
#endif 

  EXTPCKalDetector  * TrackFitterKalman::_kaldetPtr = 0;
  TrackFitterKalman * TrackFitterKalman::_theKalmanFilterFitter = 0;

  TrackFitterKalman::TrackFitterKalman()
  {
  }

  TrackFitterKalman::~TrackFitterKalman()
  {
  }

  TrackFitterBase * TrackFitterKalman::getInstance(LCParameters const * parameters)
  {

#if 0
    // try to read the fitter parameters from the LCParameters
    EVENT::FloatVec transDefocusingVals;
    parameters->getFloatVals("Kalman_TransDefocussing",transDefocusingVals);
    
    float transDefocussing;
    if ( !transDefocusingVals.empty() )
      {
	transDefocussing = transDefocusingVals[0];
      }
    else
      transDefocussing=1.;

    EVENT::FloatVec longDefocusingVals;
    parameters->getFloatVals("Kalman_LongDefocussing",longDefocusingVals);

    double longDefocussing;
    if ( !longDefocusingVals.empty() )
      {
	longDefocussing = longDefocusingVals[0];
      }
    else
      longDefocussing=1.;

    // for the diffusion we can accept the default of 0.0 if parameter not found
    double transDiffusion = parameters->getFloatVal("Kalman_TransDiffusionCoef");
    double longDiffusion  = parameters->getFloatVal("Kalman_LongDiffusionCoef");
#else

    if (! _kaldetPtr) {
      _kaldetPtr = new EXTPCKalDetector();

      ///// Very temporary -- TKalDetCradle should be constructed for all sub-tracking systems
      TKalDetCradle & lp1 = * new TKalDetCradle;
      ///// Very temporary

      lp1.Install(*_kaldetPtr);

      ///// Very temporary -- TKalDetCradle should be closed after all sub-traking systems are installed
      lp1.Close();
      ///// Very temporary
    }

#endif

    if (_theKalmanFilterFitter) {
      // reinitialise the fitter
      new ( _theKalmanFilterFitter ) TrackFitterKalman();
    } else {
      _theKalmanFilterFitter = new TrackFitterKalman();
    }

    return _theKalmanFilterFitter;
  }

  std::string TrackFitterKalman::getRevision()
  {
    return std::string( "$Rev: 0001 $ ; TrackFitterKalman ") + TrackFitterBase::getRevision();
  }

  unsigned char TrackFitterKalman::getFitterType() const
  {
    return TrackFitterBase::FitterTypes::KALMAN;
  }

  IMPL::TrackImpl * TrackFitterKalman::fitTrack(EVENT::Track const * seedTrack)
  {
    streamlog_out(DEBUG2) << "TrackFitterKalman: copying track" << std::endl;
    // create a copy of the track which is going to be fitted
    IMPL::TrackImpl * fittedTrack = LCObjectCopier::copy(seedTrack);

    streamlog_out(DEBUG2) << "TrackFitterKalman: setting parameters" << std::endl;

    //-----------------
    // Prepare KalHits
    //-----------------

    TObjArray kalhits;
    kalhits.SetOwner();

    const Double_t kmm2cm   = 0.1;

    // Prepare for a vector of gear::TPCModule pointers

    gear::TPCParameters const &theTPCParameters = marlin::Global::GEAR->getTPCParameters();
    Int_t nmodules = theTPCParameters.getNModules();

    std::vector<const gear::TPCModule *> modules;

    for (Int_t i = 0; i < nmodules; i++) {
      modules.push_back(&theTPCParameters.getModule(i));
    }

    for (TrackerHitVec::const_iterator hitIter = seedTrack->getTrackerHits().begin(); hitIter < seedTrack->getTrackerHits().end(); hitIter++) {

      //================================
      // Apply pulse height cut here !!
      //================================

#if 0
      const Int_t kCutCharge = 100;
      Double_t charge = (* hitIter)->getdEdx();

      if (charge < kCutCharge) {
	std::cerr << "" << std::endl;
	continue;
      }
#endif

      const double * pos = (* hitIter)->getPosition(); // [x, y, z]

      // convert global position suit for KalTest

      TVector3 xv(pos[0]*kmm2cm, pos[1]*kmm2cm, pos[2]*kmm2cm);

      // get the module number and the layer number

      TrackerPulse *pulse = dynamic_cast<TrackerPulse *>((* hitIter)->getRawHits()[0]);
      Int_t module = theTPCParameters.getModule(pulse->getCellID1()).getModuleID();

      // modules[module] is a target gear::TPCModule pointer
      Int_t layer  = modules[module]->getRowNumber(pulse->getCellID0());

      if (!_kaldetPtr) {

	streamlog_out(DEBUG2) << "TrackFitterKalman:  ::::: ERROR in TrackFitterKalman ::::::"  << std::endl
			      << " Failed to cast KalDet to TPCKalDet!"     << std::endl
			      << " Abort!"                                  << std::endl;
	::abort();
      }

      ////// VERY TEMPORARY ////////////////////////////////////////
      Int_t superlayer  = (module < 2 ? 0 : (module < 5 ? 1 : 2));
      Int_t layerserial = superlayer * modules[module]->getNRows() + layer;
      ////// VERY TEMPORARY ////////////////////////////////////////

      EXTPCMeasLayer * measlayerPtr
	= static_cast<EXTPCMeasLayer *>(_kaldetPtr->At(layerserial));

#if 0
      Double_t s = measlayerPtr->CalcS(xv);
      std::cerr << "TrackFitterKalman: S = " << s << std::endl;
#endif

      // if (!measlayerPtr->IsOutSide(xv)) streamlog_out(DEBUG2) << "TrackFitterKalman: The target track hit is NOT on the Measurement layer surface" << std::endl;

#if 1
      streamlog_out(DEBUG2) << "TrackFitterKalman: meas layer ID ml = " << layerserial           << std::endl
			    << " ht: (module, layer) = (" << module << ", " << layer     << ")"  << std::endl
			    << " ml: (module, layer) = (" << measlayerPtr->GetModuleID() << ", "
			    <<                               measlayerPtr->GetLayerID () << ")"  << std::endl;
#endif

      if (measlayerPtr->GetModuleID() != module || measlayerPtr->GetLayerID() != layer) {

	streamlog_out(DEBUG2) << "TrackFitterKalman:  ::::: ERROR in TrackFitterKalman ::::::"     << std::endl
			      << " meas layer ID inconsistent for ml = " << layerserial            << std::endl
			      << " ht: (module, layer) = (" << module << ", " << layer      << ")" << std::endl
			      << " ml: (module, layer) = (" << measlayerPtr->GetModuleID() << ", "
			      <<                               measlayerPtr->GetLayerID () << ")"  << std::endl
			      << " Abort!"                                                         << std::endl;
        ::abort();
      }

      Double_t   meas [2];
      Double_t   dmeas[2];
#if 1
      Int_t      side = +1;
#else
      Int_t      side = -1;
#endif
      TKalMatrix mv   = measlayerPtr->XvToMv(xv, side);
      Double_t   d    = mv(1, 0);

      meas [0] = mv(0, 0);
      meas [1] = d;

      streamlog_out(DEBUG2) << "TrackFitterKalman: xv = ("
			    << xv.X() << ", " << xv.Y() << ", " << xv.Z() <<  ")" << std::endl;

      dmeas[0]   = measlayerPtr->GetSigmaX(d);
      dmeas[1]   = measlayerPtr->GetSigmaZ(d);
      Double_t b = EXTPCKalDetector::GetBfield();
      Double_t v = EXTPCKalDetector::GetVdrift();

#if 0
      kalhits.Add(new EXTPCHit(*measlayerPtr, meas, dmeas, side, v, TVector3(pos), b));
#else // hold TrackerHit pointer
      kalhits.Add(new EXTPCHit(*measlayerPtr, meas, dmeas, side, v, static_cast<void *>(* hitIter), b));
#endif
    }

    //--
    // Sort kalhits according to r
    //--

    kalhits.Sort();

    ///// Very temporary /////

    TIter nexthittemp(&kalhits);

    std::cerr << "# of TObjArray kalhits elements = " << kalhits.GetEntries() << std::endl;

    // -----------------------------------------
    //  Eliminate double-hits in the same layer
    // -----------------------------------------

    Int_t prevlayerserial = -99999;
    TVTrackHit *prevhitPtr = 0;

    TVTrackHit *tempHitPtr = 0;
    while ((tempHitPtr = static_cast<TVTrackHit *>(nexthittemp()))) {
      const EXTPCMeasLayer &ml = static_cast<const EXTPCMeasLayer &>(tempHitPtr->GetMeasLayer());
      Int_t layer  = ml.GetLayerID();
      Int_t module = ml.GetModuleID();

      ////// VERY TEMPORARY ////////////////////////////////////////
      Int_t superlayer  = (module < 2 ? 0 : (module < 5 ? 1 : 2));
      Int_t layerserial = superlayer * modules[module]->getNRows() + layer;
      ////// VERY TEMPORARY ////////////////////////////////////////

      std::cerr << "TrackFitterKalman: Before hit elimination (module, layer) = (" << module << ", "
		<<                                                                    layer  << ")"  << std::endl;

      if (layerserial == prevlayerserial) {
	kalhits.Remove(prevhitPtr);
	kalhits.Remove(tempHitPtr);
	delete prevhitPtr;
	delete tempHitPtr;
	prevhitPtr = 0;
	tempHitPtr = 0;
	continue;
      }
      prevlayerserial = layerserial;
      prevhitPtr = tempHitPtr;

      std::cerr << "TrackFitterKalman: After hit elimination (module, layer) = (" << module << ", "
		<<                                                                   layer  << ")"  << std::endl;
    }
    kalhits.Compress();

    std::cerr << "# of TObjArray kalhits elements after the deletion = " << kalhits.GetEntries() << std::endl;
    ///// Very temporary /////

    //--
    // Set start helix parameters.
    //--

    const Int_t kMinHits = 5;

    if (kalhits.GetEntries() < kMinHits) {

      streamlog_out(DEBUG2) << "TrackFitterKalman: <<<<<< Shortage of Hits! nhits = " 
			    << kalhits.GetEntries() << " >>>>>>>" << std::endl;
      fittedTrack->setTypeBit( FITFAILEDBIT );

      return fittedTrack;
    }

    Int_t i1, i2, i3; // (i1,i2,i3) = (1st,mid,last) hit to filter

    if (gkDir == kIterBackward) {
#if 1
      i3 = 1;                        // avoid edge
      i1 = kalhits.GetEntries() - 2; // avoid edge
      i2 = i1 / 2;
#else
      i3 = 3;                        // avoid edge
      i1 = kalhits.GetEntries() - 4; // avoid edge
      i2 = i1 / 2;
#endif
    } else {
      i1 = 1;                        // avoid edge
      i3 = kalhits.GetEntries() - 2; // avoid edge
      i2 = i3 / 2;
    }

    // ----------------------------
    //  Create a dummy site: sited
    // ----------------------------

    TVTrackHit *ht1Ptr = static_cast<TVTrackHit *>(kalhits.At(i1));
    TVTrackHit *htdPtr = new EXTPCHit(*static_cast<EXTPCHit *>(ht1Ptr));
    TVTrackHit &hitd   = *htdPtr;

    hitd(0,1) = 1.e6;    // give a huge error to d
    hitd(1,1) = 1.e6;    // give a huge error to z

    TKalTrackSite &sited = *new TKalTrackSite(hitd);
    sited.SetHitOwner(); // site owns hit
    sited.SetOwner();    // site owns states

    // ----------------------
    //  Create initial helix
    // ----------------------

    TVTrackHit &h1 = *static_cast<TVTrackHit *>(kalhits.At(i1)); // first hit
    TVTrackHit &h2 = *static_cast<TVTrackHit *>(kalhits.At(i2)); // middle hit
    TVTrackHit &h3 = *static_cast<TVTrackHit *>(kalhits.At(i3)); // last hit
    TVector3    x1 = h1.GetMeasLayer().HitToXv(h1);
    TVector3    x2 = h2.GetMeasLayer().HitToXv(h2);
    TVector3    x3 = h3.GetMeasLayer().HitToXv(h3);
    THelicalTrack helstart(x1, x2, x3, h1.GetBfield(), gkDir);   // initial helix 

    // --------------------------
    //  Set dummy state to sited
    // --------------------------

    static TKalMatrix svd(kSdim,1);
    svd(0,0) = 0.;                        // dr
    svd(1,0) = helstart.GetPhi0();        // phi0
    svd(2,0) = helstart.GetKappa();       // kappa
    svd(3,0) = 0.;                        // dz
    svd(4,0) = helstart.GetTanLambda();   // tan(lambda)
    if (kSdim == 6) svd(5,0) = 0.;        // t0

    static TKalMatrix C(kSdim, kSdim);
    for (Int_t i = 0; i < kSdim; i++) {
      C(i,i) = 1.e4;   // dummy error matrix
    }

    sited.Add(new TKalTrackState(svd, C, sited, TVKalSite::kPredicted));
    sited.Add(new TKalTrackState(svd, C, sited, TVKalSite::kFiltered));

    // ---------------------------
    //  Add sited to the kaltrack
    // ---------------------------

    EXHYBTrack kaltrack;   // a track is a kal system
    kaltrack.SetOwner();   // kaltrack owns sites
    kaltrack.Add(&sited);  // add the dummy site to this track

    // ---------------------------
    //  Prepare hit iterrator
    // ---------------------------

    TIter nexthit(&kalhits, gkDir); // come in to IP, if gkDir = kIterBackward

    // ---------------------------
    //  Start Kalman Filter
    // ---------------------------

    TVTrackHit *hitPtr = 0;
    while ((hitPtr = dynamic_cast<TVTrackHit *>(nexthit()))) {
      TKalTrackSite &site = *new TKalTrackSite(*hitPtr); // new site
      if (!kaltrack.AddAndFilter(site)) {                // filter it
	streamlog_out(DEBUG2) << "TrackFitterKalman: site discarded!" << std::endl;
	delete &site;                                    // delete it if failed
#if 0
      }
#else
      } else {
	site.GetCurState().DebugPrint();
	site.DebugPrint();
	EXTPCHit             &hit = *static_cast<EXTPCHit *>(hitPtr);
        const EXTPCMeasLayer &ml = static_cast<const EXTPCMeasLayer &>(hit.GetMeasLayer());
	const TrackerHit *hp = static_cast<const TrackerHit *>(static_cast<const EXTPCHit &>(site.GetHit()).GetHitPtr());

	TVector3 xv     = ml.HitToXv(*hitPtr);
        Int_t    module = ml.GetModuleID();
        Int_t    layer  = ml.GetLayerID();
	Double_t s      = ml.CalcS(xv);
	Double_t charge = hp->getdEdx();

	std::cerr << "TrackFitterKalman: xvhit = (" << xv.X() << "," << xv.Y() << "," << xv.Z() << ")" << std::endl;
	std::cerr << "TrackFitterKalman: ml: (module, layer) = (" << module << ", "
		  <<                                                 layer  << ")"  << std::endl;
	std::cerr << "TrackFitterKalman: S = " << s                                 << std::endl;
	std::cerr << "TrackFitterKalman: Charge = " << charge                       << std::endl;
	std::cerr << "==================== Next event ====================" << std::endl;
      }
#endif
    } // end of Kalman filter

    // ---------------------------
    //  Save it as a JHelix
    // ---------------------------

    Int_t    ndf  = kaltrack.GetNDF();
    Double_t chi2 = kaltrack.GetChi2();
    Double_t cl   = TMath::Prob(chi2, ndf);
    Double_t dr   = kaltrack.GetCurSite().GetCurState()(0, 0); 
    Double_t fi0  = kaltrack.GetCurSite().GetCurState()(1, 0); 
    Double_t cpa  = kaltrack.GetCurSite().GetCurState()(2, 0); 
    Double_t dz   = kaltrack.GetCurSite().GetCurState()(3, 0); 
    Double_t tnl  = kaltrack.GetCurSite().GetCurState()(4, 0); 
#if 0
    Double_t cs   = tnl/TMath::Sqrt(1.+tnl*tnl);
    Double_t t0   = kaltrack.GetCurSite().GetCurState()(5, 0); 
#endif
    Double_t rho  = static_cast<TKalTrackState &>(kaltrack.GetCurSite().GetCurState()).GetHelix().GetRho(); 
    TVector3 xv0  = static_cast<TKalTrackSite &>(kaltrack.GetCurSite()).GetPivot();
    Double_t phi1st = 0.;
    Double_t philst = - (x3 - x1).Perp()/rho;

#if 1
    kaltrack.GetCurSite().GetCurState().DebugPrint();
    streamlog_out(DEBUG2) << "TrackFitterKalman: (chi2, ndf) = (" << chi2 << "," << ndf << ")" << std::endl;
#endif

    // ---------------------------
    //  Smooth the track
    // ---------------------------

    kaltrack.SmoothBackTo(1);

#if 1 // Very temporary
    // ====================
    //  Monitor Fit Result
    // ====================
    //
    static TFile    *file       = 0;

    static TNtupleD *hTracks    = 0;

    if (!hTracks) {
      file = new TFile("ntuple.root", "RECREATE");      
      stringstream list;
      list << "ndf:chi2:cl:dr:fi0:cpa:dz:tnl" << ends;
      hTracks = new TNtupleD("hTracks","",list.str().data());
    }
    file->cd();

    static const Int_t kNdata = 1000;
    Double_t data[kNdata];
    data[0] = ndf;
    data[1] = chi2;
    data[2] = cl;
    data[3] = dr;
    data[4] = fi0;
    data[5] = cpa;
    data[6] = dz;
    data[7] = tnl;
    hTracks->Fill(data);

    /////////// TEMPORARY TREATMENT /////////////////////////
    static const Int_t kOffset  = 11; // ndf, chi2, ...
#if 0
    static const Int_t kNlayers = 28; // # layers / module
#else
    static const Int_t kNlayers = 84; // # of layers
#endif
    /////////// TEMPORARY TREATMENT /////////////////////////

    static TNtupleD *hResXin = 0;
    if (!hResXin) {
      stringstream list;
      list << "ndf:chi2:cl:dr:fi0:cpa:dz:tnl";
      list << ":fgIsUpperHit:fgIsMidHit:fgIsLowerHit";
      for (Int_t i=0; i<kNlayers; i++) {
        list << ":dxin" << setw(2) << setfill('0') << i;
        list << ":dxot" << setw(2) << setfill('0') << i;
      }
      for (Int_t i=0; i<kNlayers; i++) {
        list << ":ch" << setw(2) << setfill('0') << i;
      }
      list << ends;
      hResXin = new TNtupleD("hResXin","",list.str().data());
    }
    for (Int_t i=kOffset; i<kNdata; i++) data[i] = 99999.;
    TIter nextlayer(&kaltrack);
    nextlayer();
    TKalTrackSite *sitep;

    fgIsMidHit   = kFALSE;
    fgIsUpperHit = kFALSE;
    fgIsLowerHit = kFALSE;
    while ((sitep = static_cast<TKalTrackSite *>(nextlayer()))) {

      const TrackerHit *hp = static_cast<const TrackerHit *>(static_cast<const EXTPCHit &>(sitep->GetHit()).GetHitPtr());
      ////// VERY TEMPORARY ////////////////////////////////////////
#if 1
      Int_t index = sitep->GetHit().GetMeasLayer().GetIndex();
#endif
      ////// VERY TEMPORARY ////////////////////////////////////////

#if 0 // Ingore boundary separation
      Int_t CenterPadOfCluster = hp->GetPadID();    
      Int_t MaxReadoutPad = 175; // for Module3 (channelmap_090324.dat)
      Int_t MinReadoutPad = 112; // for Module3 (channelmap_090324.dat)
      Int_t PadOffset     = 15;
      // select the hits which are near enough the center of readout region
      if (CenterPadOfCluster < MaxReadoutPad-PadOffset && CenterPadOfCluster > MinReadoutPad + PadOffset) {
        if      (index == fgLayerToLookAt)     fgIsMidHit   = kTRUE;
        else if (index == fgLayerToLookAt + 1) fgIsUpperHit = kTRUE;
        else if (index == fgLayerToLookAt - 1) fgIsLowerHit = kTRUE;
      }
#endif
      data[8]  = fgIsUpperHit;
      data[9]  = fgIsMidHit;
      data[10] = fgIsLowerHit;

      Double_t dxin = sitep->GetResVec()(0,0);
      Double_t dxot = 999999.;
      sitep->InvFilter();
      dxot = sitep->GetResVec(TVKalSite::kInvFiltered)(0,0);
      data[kOffset+2*index  ]        = dxin;
      data[kOffset+2*index+1]        = dxot;

#if 0
      Double_t charge = hp->GetCharge();
#else
      Double_t charge = hp->getdEdx();
#endif
      data[kOffset+index+2*kNlayers] = charge;

    }
    hResXin->Fill(data);

    static TNtupleD *hResZin = 0;
    if (!hResZin) {
      stringstream list;
      list << "ndf:chi2:cl:dr:fi0:cpa:dz:tnl";
      for (Int_t i=0; i<kNlayers; i++) {
        list << ":dz" << setw(2) << setfill('0') << i;
      }
      for (Int_t i=0; i<kNlayers; i++) {
        list << ":ch" << setw(2) << setfill('0') << i;
      }
      list << ends;
      hResZin = new TNtupleD("hResZin","",list.str().data());
    }
    for (Int_t i=kOffset; i<kNdata; i++) data[i] = 99999.;
    nextlayer.Reset();
    nextlayer();
    while ((sitep = static_cast<TKalTrackSite *>(nextlayer()))) {
#if 1
      Int_t index = sitep->GetHit().GetMeasLayer().GetIndex();
      const TrackerHit *hp = static_cast<const TrackerHit *>(static_cast<const EXTPCHit &>(sitep->GetHit()).GetHitPtr());
      Double_t charge = hp->getdEdx();
#else
      const JHit  *hp   = static_cast<const JHit *>(static_cast<const EXTPCHit &>(sitep->G
									      etHit()).GetHitPtr());
      Int_t index = hp->GetLayerID();
      Double_t charge = hp->GetCharge();
#endif

      // cerr << "index = " << index << endl;
      TKalMatrix resvec = sitep->GetResVec();
      Double_t dz = resvec(1,0);
      data[index+kOffset] = dz;
      data[index+kOffset+kNlayers] = charge;
    }
    hResZin->Fill(data);
#endif

    //--
    // put the track parameters on Track collection
    //--
    static const Int_t kMinNDF = 1;

    if (ndf >= kMinNDF) {

      fittedTrack->setD0(dr/kmm2cm);
      fittedTrack->setPhi(fi0);
      fittedTrack->setTanLambda(tnl);
      fittedTrack->setZ0(dz/kmm2cm);
      fittedTrack->setOmega(rho/kmm2cm);
      fittedTrack->setChi2(chi2);

    } else {

      streamlog_out( WARNING1 ) << "TrackFitterKalman: Fit failed, chi square is " 
				<< chi2  << std::endl;      
      fittedTrack->setTypeBit( FITFAILEDBIT );

    } // don't change the seed parameters. The ChiSquare is zero, so these tracks can be 
      // identified.

    // finally set the fitter type in the trackType word
    setTrackFitterType( this->getFitterType(), fittedTrack );

    return fittedTrack;
  }

} // namespace marlintpc