#include "TrackingMods/IOTrackModule.hh"

#include "AbsEnv/AbsEnv.hh"
#include "TrackingCT/CT_ReFit.hh"
#include "TrackingObjects/Storable/CT_HitSet.hh"
#include "TrackingObjects/Tracks/CdfTrack.hh"
#include "TrackingObjects/Storable/CdfTrackView.hh"
#include "TrackingObjects/Tracks/track_cut.hh"
#include "TrackingObjects/CT_Track/CT_Track.hh" 
#include "TrackingObjects/CT_Track/Cdf_to_CT_Track.hh"
#include "TrackingSI/TrackFitting/SiKalmanFitter.hh"
#include "linalg/CT_Helix.hh"
#include "linalg/CT_Utils.hh"
#include "TrackingObjects/CT_Track/CT_Residual.hh"
#include "CotGeometry/CT_Configuration.hh"
#include "TrackingObjects/CT_Simulation/CT_Simulation.hh"
#include "TrackingObjects/SiData/SiCluster.hh"
#include "TrackingCT/CT_DedxCalculator.hh"
#include "TrackingObjects/Tracks/dEdxCOT.hh"

#include <cmath>
#include <set>
#include <vector>
#include <map>
//============================
//  Constructor
//============================

IOTrackModule::IOTrackModule( const char* const theName, 
                                        const char* const theDescription )
  : AppModule( theName, theDescription ),
    _selectInputTrackColl("selectInputTrackColl",this,false),
    _inputTrackDescr("inputTrackDescr",this,""),
    _processName("ProcessName",this,"PROD"),
    _writeOut("writeTracks",this,true),
    _outputTrackDescr("outputTrackDescr",this,"IOTracks"),
    _setRoadWidthFactor("setRoadWidthFactor",this,1.0),
    _insideOutFit("InsideOutFit",this,true),
    // _useFixedSizeRoad("useFixedSizeRoad",this,false),
    // _setSearchRoad("setSearchRoad",this,0.1),
    _debugOutput("debugOutput",this,false),
    _minAx("MinAxialHit", this, 6),
    _minSt("MinStereoHit", this, 2),
    _IOmode("IOtrackingMode", this, 1),
    _BCfit("2dConstraint", this, true),
    _seedCurvature("SeedCurvature",this,false),
    _writeDEDx("WriteDEDx", this, true),
    _kalFitter( 0 )
{
  // Add Commands to the menu
  commands()->append( &_selectInputTrackColl );
  commands()->append( &_inputTrackDescr );
  commands()->append( &_processName );
  commands()->append( &_writeOut );
  commands()->append( &_outputTrackDescr );
  commands()->append( &_insideOutFit );
  commands()->append( &_setRoadWidthFactor );
  commands()->append( &_debugOutput );
  commands()->append( &_minAx);
  commands()->append( &_minSt);
  commands()->append( &_IOmode);
  commands()->append( &_BCfit);
  commands()->append( &_seedCurvature);
  commands()->append( &_writeDEDx );

  // Add Some description to the commands.
  _IOmode.addDescription("\tSeed track mode.  1:  Default stand-alone \t 2: Si PT \t 3: silicon-forward \t 4: stand-alone || silicon-forward");
  _selectInputTrackColl.addDescription(
     "\tUse specified CdfTrackColl as input. If false, use silicon\n\
\tstand-alone tracks found by allTracks");
  _inputTrackDescr.addDescription(
     "\tDescription string for input CdfTrackColl"); 
  _insideOutFit.addDescription(
     "\tPerform inside-out fit on input Si SA tracks before feeding to COT"); 
  _writeOut.addDescription("\tWrite output CdfTrackColl to event");
  _outputTrackDescr.addDescription(
     "\tDescription string for output CdfTrackColl");
  _seedCurvature.addDescription("\tUse input track curvature in the COT track fit.");
  _setRoadWidthFactor.addDescription(
     "\tSet scale factor for radius dependent road width.");

  _debugOutput.addDescription("\tPrint debugging information during execution");
  _processName.addDescription("\tProcess name for input objects");
}

//============================
//  Destructor
//============================
IOTrackModule::~IOTrackModule() {
}

//=====================
//  Operations
//=====================

AppResult IOTrackModule::beginJob( AbsEvent* aJob )   {
  _kalFitter = SiKalmanFitter::instance();
  _ku = KalUtils::instance();
  return AppResult::OK;
}

AppResult IOTrackModule::beginRun( AbsEvent* aRun )   {
  if (_kalFitter) _kalFitter->refreshDetector();
  return AppResult::OK;
}

AppResult IOTrackModule::event( AbsEvent* anEvent )  
{
  CdfTrackView * inputView = 0;
  if ( _selectInputTrackColl.value() )
  {
    // Get tracks from the specified track collection
    // Filter out any tracks that already have COT hits.
    //
    CdfTrackColl_ch hTracks;
    StorableObject::ANDSelector sel = 
      StorableObject::SelectByClassName("CdfTrackColl")
      &&
      StorableObject::SelectByDescription(_inputTrackDescr.value())
      &&
      StorableObject::SelectByProcessName(_processName.value());
    EventRecord::ConstIterator iTracks( anEvent, sel );
    if ( iTracks.is_valid() ) {
	inputView = new CdfTrackView;
	CdfTrackView::collection_type &  inputTracks = inputView->contents();
	hTracks = iTracks;
	const CdfTrackColl::collection_type & tracks = hTracks->contents();
	for ( CdfTrackColl::const_iterator i = tracks.begin(), iend = tracks.end(); i != iend ; ++i ) {
	  if ( (**i).numCTHits() == 0 ) 
	    inputTracks.push_back( CdfTrackView::value_type( hTracks, i ) );
	}
	
	if (verbose()) {
	  cout << "IOTrackModule:  CdfTrackColl '" << _inputTrackDescr.value()
	       << "' with " << inputTracks.size() << " non-COT tracks" << endl;
	}
    }
  }
  else {
    // access existing tracks
    CdfTrackView_h atView;
    CdfTrackView::allTracks(atView,_processName.value());
    CdfTrackView::CollType tracks = atView->contents();
    
    // create silicon stand-alone track view
    CdfTrackView* siliconStandAloneTracks;
    
    if (_IOmode.value()==1) 
      siliconStandAloneTracks = new CdfTrackView (tracks.begin(), tracks.end(), 
						  track_cut::SelectAlgorithm(CdfTrack::KalSvxStandAloneAlg));
    else if (_IOmode.value()==2) 
      siliconStandAloneTracks = new CdfTrackView (tracks.begin(), tracks.end(), 
						  track_cut::SelectAlgorithm(CdfTrack::SiPT));
    else if (_IOmode.value()==3)
      siliconStandAloneTracks = new CdfTrackView (tracks.begin(), tracks.end(), 
						  track_cut::SelectAlgorithm(CdfTrack::KalForwardStandAloneAlg));

    else if (_IOmode.value()==4)
      siliconStandAloneTracks = new CdfTrackView (tracks.begin(), tracks.end(), 
						  track_cut::SelectAlgorithm(CdfTrack::KalSvxStandAloneAlg) || 
						  track_cut::SelectAlgorithm(CdfTrack::KalForwardStandAloneAlg));
    
    inputView = siliconStandAloneTracks;
  
    if (verbose())
      cout << "IOTrackModule:  using allTracks, " 
	   << inputView->contents().size() << " non-COT tracks" << endl;
    if (!_selectInputTrackColl.value() && _IOmode.value()==2 && verbose())
      cout << "Chose to have SiPT as IO Seed tracks\n";
  }
  
  if ( !inputView ) 
    return AppResult::OK;
  
  // Prepare output track collection
  //
  ConstHandle<SiHitSet> siliconHitSet_h;
  EventRecord::ConstIterator siliconHitSetIter(anEvent,"SiHitSet","*",_processName.value());
  if (siliconHitSetIter.is_valid()) {
    siliconHitSet_h = siliconHitSetIter;
  }  
  CdfTrackColl_h  outputColl = CdfTrackColl_h(new CdfTrackColl(siliconHitSet_h));
  outputColl->set_description( _outputTrackDescr.value() );
  CdfTrackColl::collection_type & outputTracks = outputColl->contents();
  
  // Configure fitter, initialize track loop
  //
  const CT_Configuration& cfg = CT_Configuration::reference();
  const int nsl = cfg.numSuperLayers();

  CT_ReFit refit1;
  refit1.setRoadWidthFactor( _setRoadWidthFactor.value() );

  // loop over input tracks
  CdfTrackView::iterator t, tend;
  for (t =  inputView->contents().begin(), tend = inputView->contents().end(); t != tend; ++t ) {
    
    const CdfTrack& trkRef = **t;
    CT_Track trk = Cdf_to_CT_Track(trkRef);
    
    // first, drop all the COT hits on SiPT tracks if SiPT is chosen to be the seed
    if (_IOmode.value()==2){
      for ( int s=0, ilr=0; s<cfg.numSuperLayers(); ++s) {
	const CT_HitLink* phit = trk.beginCTHitBlock(s);
	for ( int w=0; w<cfg.numWires(s); ++w,++ilr) 
	  if (phit[w]) trk.dropHit(ilr);
      }
    }
    
    if ( _insideOutFit.value() ) {
      HepVector parameters(5);
      HepSymMatrix covariance(5);
      double chi2;
      int degreesoffreedom;
      bool interactCOT;
      CdfTrack* track = new CdfTrack( trkRef );
      track->setAlgorithm(CdfTrack::KalSvxStandAloneAlg);
      bool success = _ku->forwardFit2COT(track,covariance,parameters,chi2,degreesoffreedom,interactCOT);
      if (success) trk.setFit3d(5, parameters, covariance, chi2);	    
      track->destroy();
    }

    refit1.searchNewTrack();
    refit1.restore( trk );

    if (!refit1.ok()) {
      if ( verbose() )
	cout << "IOTrackModule:  COT fit of track " << t - inputView->contents().begin() << " failed." << endl;
      continue;
    } 

    // add hits to track
    for (int sl = 0,ilr=0; sl < nsl ; ++sl ) { 
      for (int wire = 0; wire < cfg.numWires(sl); ++wire, ++ilr) {
	const CT_ReFit::LayerData& ldat = *(refit1.layer_data() + ilr); 
	if (! ldat.link.is_valid()) continue;
	if (ilr > trkRef.lastLayerCT()) continue;
	trk.addHit(sl,ldat.link);         
	int hitnumber = ldat.link.offset();
	//int particle =-999;
	//if (hitnumber < CT_Simulation::reference().hits().size()) 
	//  particle =CT_Simulation::reference().hits()[hitnumber].parent();
	
	//if (_debugOutput.value())
	//  std::cout << "ilr: " << ilr << ", ldat.deltaD(): " << ldat.deltaD() << ", parent=" << particle << endl;
      }
    }

    if (_debugOutput.value()) {
      std::cout << "  Before COT Fit:" << std::endl
		<< "    cot = " << trk.par()[0] << " +\\- "
		<< sqrt( trk.cov()[0][0] ) << std::endl
		<< "    cur = " << std::setprecision(8) << trk.par()[1] << " +\\- " 
		<< sqrt( trk.cov()[1][1] ) << std::endl
		<< "    z0  = " << trk.par()[2] << " +\\- " 
		<< sqrt( trk.cov()[2][2] ) << std::endl
		<< "    d0  = " << trk.par()[3] << " +\\- " 
		<< sqrt( trk.cov()[3][3] ) << std::endl
		<< "    p0  = " << trk.par()[4]  << " +\\- " 
		<< sqrt( trk.cov()[4][4] ) << std::endl
		<< "    N COT hits: " << trk.numCTHits() 
		<< "    Last COT layer: " << trkRef.lastLayerCT() << std::endl;
    }
	  
    // now fit the track
    CT_ReFit refit;

    if (!_seedCurvature.value()) {

      refit(trk,_minAx.value(), _minSt.value());

      if (!refit.ok()) {
	if ( verbose() )
	  std::cout << "IOTrackModule: COT fit of restored track " << 
	    t - inputView->contents().begin() << 
	    " failed." << std::endl;
	continue;
      } 

    } else {

      // use seed track curvature in the fit
      int status = -2;
      int iter = 0;
      HepSymMatrix weight(5);
      HepVector gradient(5);
      HepVector dpars(5);
      HepVector pars(5);
      float dchi2 = 9999.;

      while (++iter < 10 && status == -2) {

	bool initialize = (iter==1);
	refit.prepareRefitDicosmic(trk,initialize);

	weight = refit.wgt();
	gradient = refit.grad();
	pars = refit.par();
	
	// add track curvature to gradient and weight
	gradient[CT_Standard::icu] += (pars(CT_Standard::icu) - trk.curvature()) /
	  trk.cov()[CT_Standard::icu][CT_Standard::icu];
	
	weight[CT_Standard::icu][CT_Standard::icu] += 1./trk.cov()[CT_Standard::icu][CT_Standard::icu];
	
	// invert weight matrix to get covariance
	int ier = -1;
	weight.invert(ier);
	if (ier!=0) status = -1; // failed fit

	// change in paramters
	dpars = weight*gradient;
	pars -= dpars;
	
	dchi2 = dot(gradient,dpars);
	if (dchi2 < 0.0) status = 0; // diverged
	else if (dchi2 < 1.) status = 1; // converged

      }  // end of iterations

      if (status!=1) {
	if ( verbose() )
	  std::cout <<"IOTrackModule: COT fit of restored track " << 
	    t - inputView->contents().begin() << " failed." << std::endl;
	continue;
      } 

    }

    // compute chi2
    double chi2 = 0.;
    for (int sl = 0, ilr=0; sl < nsl ; ++sl ) { 
      for (int wire = 0; wire < cfg.numWires(sl); ++wire, ++ilr) {
	const CT_ReFit::LayerData& ldat = *(refit.layer_data() + ilr); 
	if (! ldat.link.is_valid()) continue;
	chi2 += square(ldat.deltaD()) * ldat.weight;
	if (_debugOutput.value()) std::cout << "ilr: " << ilr << ", ldat.deltaD(): " << ldat.deltaD() << endl;
      }
    }
    if (_debugOutput.value()) std::cout << "Before BC, chi2=" << chi2 << endl;

    trk.setFit3d(refit.npar(), refit.par(), refit.cov(), chi2);
    if (_BCfit.value()) {
      // constraint trk d0 and z0 to SiSA
      double d0=trkRef.par()[3], p0=trkRef.par()[4], d0e=sqrt(trkRef.cov()[3][3]);
      double z0=trkRef.par()[2], z0e=sqrt(trkRef.cov()[2][2]);
      userBeam.setIsValid(true);
      userBeam.setBeamX(-1.0*sin(p0)*d0); 
      userBeam.setBeamY(     cos(p0)*d0);
      userBeam.setWidthX(fabs(sin(p0)*d0e));
      userBeam.setWidthY(fabs(cos(p0)*d0e));
      refit.setZVertex(z0,z0e);
      refit.initBeamPos(userBeam);
      refit.beamcon3d(trk);
      
      if (_debugOutput.value()) 
	std::cout << "  Before BC Fit:" << std::endl
		  << "    cot = " << trk.par()[0] << " +\\- "
		  << sqrt( trk.cov()[0][0] ) << std::endl
		  << "    cur = " << std::setprecision(8) << trk.par()[1] << " +\\- " 
		  << sqrt( trk.cov()[1][1] ) << std::endl
		  << "    z0  = " << trk.par()[2] << " +\\- " 
		  << sqrt( trk.cov()[2][2] ) << std::endl
		  << "    d0  = " << trk.par()[3] << " +\\- " 
		  << sqrt( trk.cov()[3][3] ) << std::endl
		  << "    p0  = " << trk.par()[4]  << " +\\- " 
		  << sqrt( trk.cov()[4][4] ) << std::endl
		  << "    N COT hits: " << trk.numCTHits() << std::endl;
      
      if (!refit.ok()) {
	if ( verbose() )
	  cout << "IOTrackModule: BC of COT restored track " << t - inputView->contents().begin() << " failed." << endl;
	continue;
      }

      trk.setFit3d(refit.npar(), refit.par(), refit.cov(), chi2);

      CT_ReFit refit2;  // to get the right chi2
      refit2.initLayerData(trk);
      chi2=0;
      for ( int sl = 0, ilr=0; sl < nsl ; ++sl ) {
        for (int wire = 0; wire < cfg.numWires(sl); ++wire, ++ilr) {
          const CT_ReFit::LayerData& ldat = *(refit2.layer_data() + ilr);
          if (! ldat.link.is_valid()) continue;
          chi2 += square(ldat.deltaD()) * ldat.weight;
	}
      }
      
    }
    if (_debugOutput.value()) 
      std::cout << "  After the last COT refit:" << std::endl
		<< "    cot = " << trk.par()[0] << " +\\- "
		<< sqrt( trk.cov()[0][0] ) << std::endl
		<< "    cur = " << std::setprecision(8) << trk.par()[1] << " +\\- " 
		<< sqrt( trk.cov()[1][1] ) << std::endl
		<< "    z0  = " << trk.par()[2] << " +\\- " 
		<< sqrt( trk.cov()[2][2] ) << std::endl
		<< "    d0  = " << trk.par()[3] << " +\\- " 
		<< sqrt( trk.cov()[3][3] ) << std::endl
		<< "    p0  = " << trk.par()[4]  << " +\\- " 
		<< sqrt( trk.cov()[4][4] ) << std::endl
		<< "    N COT hits: " << trk.numCTHits() << std::endl;

    if ( _writeOut.value() ) {
      CdfTrack* track = new CdfTrack( trkRef );
      for ( int sl = 0, ilr=0; sl < nsl ; ++sl ) { 
	const CT_HitLink* phit = trk.beginCTHitBlock(sl);
	for (int wire = 0; wire < cfg.numWires(sl); ++wire,++ilr) {
	  track->dropHit(ilr);
	  if ( !phit[wire] ) continue;
	  track->addHit(sl,phit[wire]);         		   
	}
      }

      track->setFit(HelixFit(chi2, (trk.numCTHits()-refit.npar()), refit.par(),refit.cov()));
      CdfTrack* parTrk = new CdfTrack();
      for ( int sl = 0, ilr=0; sl < nsl ; ++sl ) { 
	const CT_HitLink* phit = trk.beginCTHitBlock(sl);
	const CT_Residual* res = trk.beginCTResiduals(sl);
	for (int wire = 0; wire < cfg.numWires(sl); ++wire,++ilr) {
	  if ( !phit[wire] ) continue;
	  parTrk->addHit(sl,phit[wire]);
	  parTrk->setResidual(ilr,res[wire]);
	  if (trk.obscuredCT().test(ilr)) parTrk->markObscured(ilr);
	}
      }
      parTrk->setFit(HelixFit(chi2, (trk.numCTHits()-refit.npar()), refit.par(),refit.cov()));
      parTrk->setChi2CT( trk.chi2() );
      parTrk->setMatch( trk.match() );
      if (trk.isDuplicate()) parTrk->markDuplicate();
      parTrk->setAlgorithm(CdfTrack::CotStandAloneAlg);

      if (_writeDEDx.value()) { 
	dEdxCOT dedx;
	CT_DedxCalculator(*parTrk, dedx);
	parTrk->setDedxCT( dedx.COTOfflineCorrDedx(), dedx.NumOfflineCorrHits() );
      }

      parTrk->getAlgorithm().SetCotParentAlgo(CdfTrack::TrackingAlgorithm::SiParent);

      const CdfTrack_cfwdlnk oldparent(parTrk);
      track->setParentLink(oldparent);
      track->setAlgorithm(CdfTrack::OutsideIn3DAlg);

      if(_kalFitter->fit( *track )!= TrackFitter::OK) {
	if (verbose()) cout << "**** _kalFitter failed\n";
	track->destroy();
	parTrk->destroy();
	continue;
      }

      // setting parent to COT, child to SI
      // will remove SI in PadTrackMaker and keep COT
      track->setChildLink(CdfTrack_cfwdlnk(&trkRef));
      parTrk->setChildLink(CdfTrack_cfwdlnk(track));

      track->setAlgorithm(CdfTrack::InsideOutAlg);
      outputTracks.push_back(track);
      outputTracks.push_back(parTrk);


      if ( _debugOutput.value() )  {
	std::cout << "IOTrackModule:  refit OK" << std::endl
		  << "  Before IO:" << std::endl
		  << "    cot = " << trkRef.par()[0] << " +\\- "
		  << sqrt( trkRef.cov()[0][0] ) << std::endl
		  << "    cur = " << trkRef.par()[1] << " +\\- " 
		  << sqrt( trkRef.cov()[1][1] ) << std::endl
		  << "    z0  = " << trkRef.par()[2] << " +\\- " 
		  << sqrt( trkRef.cov()[2][2] ) << std::endl
		  << "    d0  = " << trkRef.par()[3] << " +\\- " 
		  << sqrt( trkRef.cov()[3][3] ) << std::endl
		  << "    p0  = " << trkRef.par()[4]  << " +\\- " 
		  << sqrt( trkRef.cov()[4][4] ) << std::endl
		  << "    N COT hits: " << trkRef.numCTHits() << std::endl
		  << "    N SI  hits: " << trkRef.numSIHits() << std::endl;
	  
	std::cout << "IOTrackModule:  refit OK" << std::endl
		  << "  After IO:" << std::endl
		  << "    cot = " << track->par()[0] << " +\\- "
		  << sqrt( track->cov()[0][0] ) << std::endl
		  << "    cur = " << track->par()[1] << " +\\- " 
		  << sqrt( track->cov()[1][1] ) << std::endl
		  << "    z0  = " << track->par()[2] << " +\\- " 
		  << sqrt( track->cov()[2][2] ) << std::endl
		  << "    d0  = " << track->par()[3] << " +\\- " 
		  << sqrt( track->cov()[3][3] ) << std::endl
		  << "    p0  = " << track->par()[4]  << "+\\- " 
		  << sqrt( track->cov()[4][4] ) << std::endl
		  << "    N COT hits: " << track->numCTHits() << std::endl
		  << "    N SI  hits: " << track->numSIHits() << std::endl;
	
      }
    } // if write      
  } // track loop
  
  ConstLink<CdfTrackColl> collLink( outputColl );
  for (CdfTrackColl::const_iterator i=outputColl->contents().begin();i!=outputColl->contents().end();++i) {
    if ( (**i).child().is_nonnull() ) 
      (**i).child()->setParentLink( CdfTrack_cfwdlnk(collLink, i));    
  }

  if ( _writeOut.value() ) {
    outputColl->assignTrackIds(anEvent);
    if ( currentRCPID().isValid() ) outputColl->set_rcp_id( currentRCPID() );
    anEvent->append( outputColl );
  }  
  // Free up memory
  inputView->destroy();
  return AppResult::OK;
}

AppResult IOTrackModule::endRun( AbsEvent* aRun )   {
  return AppResult::OK;
}

AppResult IOTrackModule::endJob( AbsEvent* aJob )   {
  return AppResult::OK;
}

AppResult IOTrackModule::abortJob( AbsEvent* aJob )   {
  return AppResult::OK;
}