// PTrack.cpp

#include "PTrack.h"
#include <iostream>
#include "trfutil/trfstream.h"
#include "Path.h"
#include "trflayer/Layer.h"
#include "trffit/AddFitter.h"
#include "trffit/FullFitter.h"
#include "Checker.h"
#include "trfobj/TrfReportingObject.h"
#include "trfbase/Miss.h"
#include "trflayer/LTrack.h"
#include "trfutil/TupleManager.h"
#include "trfutil/EventID.h"
#include "trfutil/CpuTimerTuple.h"
#include "trfcut/CutRecord.hpp"
#include "trfcut/CutRecorder.hpp"
#include "trflayer/ClusterContainer.h"

using std::ostream;
using std::cout;

using trf::Surface;
using trf::ETrack;
using trf::Layer;
using trf::Propagator;
using trf::CandidateLayer;
using trf::HitPtr;
using trf::HitList;
using trf::MissPtr;
using trf::ClusterFilter;
using trf::CutRecord;
using trf::CutRecorder;
using trf::MTrack;
using trf::LayerStatChain;
using trf::LTrackList;
using trf::Path;
using trf::PTrack;

//************************************************************************
// Static attributes.
//************************************************************************

int PTrack::_count = 0;

bool PTrack::_ltuple = false;

//************************************************************************
// Private methods.
//************************************************************************

// complete constructor

PTrack::PTrack(const Path& path, const MTrack& trm, 
               const LayerStatChain& chain, Tuple* ptup, CutRecord* prec)
: RefCounter(), _path(path), _trm(trm), _chain(chain), _stopped(false),
  _end_of_path(false), _ptup(ptup),
  _id(++_count), _parent_id(-1),
  _pcut_record(prec) {
  // Assign the cut ID's.
  if ( _pcut_record != 0 ) {
    _pcut_record->set_ptrack_id(get_id());
    _pcut_record->set_path_id(get_path().get_id());
  } else {
    assign_cut_record();
  }
}

//************************************************************************

// output stream

void PTrack::ostr(ostream& stream) {
  stream << begin_object;
  stream << "PTrack id: " << get_id() << new_line;
  stream << "Parent id: " << get_parent_id() << new_line;
  stream <<  _trm << new_line;
  stream << _path << new_line;
  stream << _chain << new_line;
  stream << "PTrack is ";
  if ( ! is_stopped() ) stream << "not ";
  stream << "stopped." << new_line;
  if ( _ptup ) 
    stream << "Ntuple name is " << _ptup->title() << ".";
  else 
    stream << "Ntuple is not defined.";
  stream << new_line;
  if ( _pcut_record)
    stream << *_pcut_record;
  stream << end_object;
}

//************************************************************************

// Check the track: return true if all checkers pass.

bool PTrack::check() {
  const Path::CheckerList& checkers = _path.get_checkers();
  for ( Path::CheckerList::const_iterator ichk=checkers.begin(); 
        ichk!=checkers.end(); ++ichk )
    if ( ! (**ichk).status_with_record(_trm,_pcut_record) ) return false;
  return true;
}

//************************************************************************

// Post check the track: return true if all post checkers pass.

bool PTrack::post_check() {
  const Path::CheckerList& checkers = _path.get_post_checkers();
  for ( Path::CheckerList::const_iterator ichk=checkers.begin(); 
        ichk!=checkers.end(); ++ichk )
    if ( ! (**ichk).status_with_record(_trm,_pcut_record) ) return false;
  return true;
}

//************************************************************************

// Assign the cut record.

void PTrack::assign_cut_record() {
  assert( _pcut_record == 0 );
  CutRecorder* precorder = CutRecorder::instance();
  if ( precorder!=0 && precorder->enabled() ) {
    _pcut_record = precorder->
      make_cut_record(-1, get_id(), get_path().get_id());
  }
}

//************************************************************************
// Public methods.
//************************************************************************

// constructor from kinematic track and head path

PTrack::PTrack(const Path& path, const ETrack& tre)
: RefCounter(),
  _path(path), _trm(tre), _chain(), _stopped(false), _end_of_path(false),
  _ptup(0), _id(++_count), _parent_id(-2),
  _pcut_record(0) {
  assign_cut_record();
}

//************************************************************************

// constructor from hit/miss track and head path

PTrack::PTrack(const Path& path, const MTrack& trm)
: RefCounter(),
  _path(path), _trm(trm), _chain(), _stopped(false), _end_of_path(false),
  _ptup(0), _id(++_count), _parent_id(-3),
  _pcut_record(0) {
  assign_cut_record();
}

//************************************************************************

// destructor

PTrack::~PTrack() { }

//************************************************************************

// Propagate this track to the next surface in this layer or the first
// surface in each next layer.
// Return all such tracks.

PTrack::PTrackList 
PTrack::propagate(bool debug, CpuTimerTuple* ptimer,
                  const TrfReportingObject* prep) const {

  // If a timer is passed, then assign names for the stages to
  // be timed.
  string total_time = "";
  string propagate_time = "";
  string fetch_cluster_time ="";
  string filter_cluster_time ="";
  string fit_time = "";
  string original_time = "";
  string post_time = "";
  if ( ptimer != 0 ) {
    string name = _path.get_name();
    bool subclocks = false;
    if ( name.size() > 0 ) {
      char chlast = name[name.size()-1];
      if ( chlast == '*' ) {
        string::iterator inam=name.end();
        --inam;
        name.erase(inam);
        subclocks = true;
      }
    }
    total_time = "pt" + name;
    if ( ! ptimer->is_defined(total_time) ) {
      ptimer->add_subclock(total_time);
    }
    if ( subclocks ) {
      propagate_time = "pp" + name;
      fetch_cluster_time = "pc" + name;
      filter_cluster_time = "ps" + name;
      fit_time = "pf" + name;
      original_time = "po" + name;
      post_time = "pa" + name;
      if ( ! ptimer->is_defined(propagate_time) ) {
        ptimer->add_subclock(propagate_time);
      }
      if ( ! ptimer->is_defined(fetch_cluster_time) ) {
        ptimer->add_subclock(fetch_cluster_time);
      }
      if ( ! ptimer->is_defined(filter_cluster_time) ) {
        ptimer->add_subclock(filter_cluster_time);
      }
      if ( ! ptimer->is_defined(fit_time) ) {
        ptimer->add_subclock(fit_time);
      }
      if ( ! ptimer->is_defined(original_time) ) {
        ptimer->add_subclock(original_time);
      }
      if ( ! ptimer->is_defined(post_time) ) {
        ptimer->add_subclock(post_time);
      }
    }
  }

  // Start the total clock.
  if ( total_time != "" ) {
    int stat = ptimer->start(total_time);
    assert( stat == 0 );
  }

  // Create the initial list of new tracks.
  PTrackList new_tracks;

  // Do not propagate a stopped track.
  assert( ! is_stopped() );
  if ( is_stopped() ) return new_tracks;

  // Counter for the number of child paths
  int ichildpath = 0;

  // Counter for new propagated tracks.
  int iprop = 0;

  // Counters for new hits.
  int iaddhit_attempt = 0;
  int iaddhit_success = 0;

  // counter for new misses
  int iaddmiss_attempt = 0;
  int iaddmiss_success = 0;

  // If this is the head path or the end of a candidate layer, then
  // construct a new track for each child path and propagate each such
  // track to the first surface in its candidate layer.
  if ( propagate_time != "" ) {
    int stat = ptimer->start(propagate_time);
    assert( stat == 0 );
  }
  if (  _path.is_head()
//        || !_chain.current_status_is_valid() 
        ||  _chain.get_current_status().at_exit()  ) {
    // loop over child paths
    Path::PathList children = _path.get_children();
    Path::PathList::const_iterator ipth;
    for ( ipth=children.begin(); ipth!=children.end(); ++ipth ) {
      ++ichildpath;
      const Path& path = **ipth;
      CandidateLayer cnl = path.get_candidate_layer();
      assert( cnl.is_valid() );
      const Layer& lyr = cnl.get_layer();
      const Propagator& prop = cnl.get_propagator();
      // propagate to the next candidate
      LTrackList ltracks;
      lyr.propagate(get_track().get_track(), prop, ltracks);
      // Create a PTrack for each propagated track.
      LTrackList::const_iterator itrl;
      for ( itrl=ltracks.begin(); itrl!=ltracks.end(); ++itrl ) {
        ++iprop;
        const LTrack& trl = *itrl;
        assert( trl.at_top_status() );
        const LayerStatChain& chain = trl.get_status_chain();
        const ETrack& tre = trl.get_track();
        MTrack trm = _trm;
        trm.set_fit( tre, trm.get_chi_square() );
	CutRecord* prec = 0;
	if(_pcut_record)
	  prec = _pcut_record->make_cut_record();
        new_tracks.push_back(
          PTrackPtr( new PTrack(path,trm,chain,get_tuple(), prec) ) );
        new_tracks.back()->set_parent_id( get_id() );
      }  // end loop over propagated tracks
    }  // end loop over candidate layers
  }  // end propagation to new candidates

  // If we are inside a layer, copy the track and propagate the copy
  // to the next surface in its candidate layer.
  else {
    assert( ! _chain.get_current_status().at_exit() );
    const CandidateLayer& cnl = _path.get_candidate_layer();
    assert( cnl.is_valid() );
    const Layer& lyr = cnl.get_layer();
    const Propagator& prop = cnl.get_propagator();
    // propagate to the next surface
    LTrack trl;
    trl.get_track() = _trm.get_track();
    trl.get_status_chain() = _chain;
    LTrackList ltracks;
    lyr.propagate(trl,prop,ltracks);
    // Create a PTrack for each propagated track.
    LTrackList::const_iterator itrl;
    for ( itrl=ltracks.begin(); itrl!=ltracks.end(); ++itrl ) {
      ++iprop;
      const LTrack& trl = *itrl;
      assert( trl.at_top_status() );
      const LayerStatChain& chain = trl.get_status_chain();
      const ETrack& tre = trl.get_track();
      MTrack trm = _trm;
      trm.set_fit( tre, trm.get_chi_square() );
      CutRecord* prec = 0;
      if(_pcut_record)
	prec = _pcut_record->make_cut_record();
      new_tracks.push_back( 
        PTrackPtr( new PTrack(_path,trm,chain,get_tuple(), prec) ) );
      new_tracks.back()->set_parent_id( get_id() );
    }  // end loop over propagated tracks
  }  // end propagation within current layer
  if ( propagate_time != "" ) {
    int stat = ptimer->stop(propagate_time);
    assert( stat == 0 );
  }

  if ( prep != 0 ) {
    prep->make_report(*this, "PTrack propagated", new_tracks);
  }

  // Loop over the new tracks.
  // One miss and all nearby hits are added to each new track.
  // The list may be modified but we iterate over the original list.
  PTrackIterator itrp = new_tracks.begin();
  int size = new_tracks.size();
  for ( int icnt=0; icnt<size; ++icnt ) {
    PTrack& trp = **itrp;
    const Path& path = trp._path;
    const LayerStatChain& chain = trp._chain;
    const CandidateLayer& cnl = path.get_candidate_layer();
    assert( cnl.is_valid() );
    const AddFitter& fit = cnl.get_fitter();
    CutRecord* pcut_record1 = trp._pcut_record;
  
    // Fetch the nearby clusters for the new track.
    if ( fetch_cluster_time != "" ) {
      int stat = ptimer->start(fetch_cluster_time);
      assert( stat == 0 );
    }
    const ETrack& tre = trp._trm.get_track();
    ClusterFilter::CutRecordMap* precs = 0;
    if ( pcut_record1 != 0 )
      precs = new ClusterFilter::CutRecordMap;
    const ClusterContainer* pbox = 
      &chain.get_clusters(tre, pcut_record1, precs);
    if ( fetch_cluster_time != "" ) {
      int stat = ptimer->stop(fetch_cluster_time);
      assert( stat == 0 );
    }

    // If the cluster finder has not implemented cut recording, build
    // the cut record map here.
    // We will have to assign the PTrack and Path ID's for the new cut
    // record if and when it is put on a PTrack.
    {
      const ClusterList& clusters = pbox->get_clusters();
      if ( pcut_record1 != 0 && precs->size() == 0 && clusters.size() > 0 ) {
        ClusterFilter::CutRecordMap& recs = *precs;
        for ( ClusterList::const_iterator iclu=clusters.begin();
	      iclu!=clusters.end(); ++iclu ) {
          recs[*iclu] = pcut_record1->make_cut_record();
        }
      }
    }

    // Report nearby clusters.
    if ( prep != 0 ) {
      prep->make_report(*this, "PTrack nearby clusters", *pbox);
    }
    if ( filter_cluster_time != "" ) {
      int stat = ptimer->start(filter_cluster_time);
      assert( stat == 0 );
    }
    const Path::ClusterFilterList& cfilters=path.get_cluster_filters();
    for ( Path::ClusterFilterList::const_iterator ifil=cfilters.begin();
          ifil!=cfilters.end(); ++ifil ) {
      pbox = &(*ifil)->filter_container(trp._trm, *pbox, precs);
      // Report filtered clusters.
      if ( prep != 0 ) {
        prep->make_report(*this, "PTrack cluster filter", **ifil);
        prep->make_report(*this, "PTrack cluster filter result", *pbox);
      }
    }
    // Report the cluster cut record map.
    if ( prep != 0 ) {
      prep->make_report(*this, "PTrack cluster cut records", precs);
    }
    if ( filter_cluster_time != "" ) {
      int stat = ptimer->stop(filter_cluster_time);
      assert( stat == 0 );
    }

    // Loop over clusters.
    if ( fit_time != "" ) {
      int stat = ptimer->start(fit_time);
      assert( stat == 0 );
    }
    const ClusterList& clusters = pbox->get_clusters();
    for ( ClusterList::const_iterator iclu=clusters.begin();
          iclu!=clusters.end(); ++iclu ) {
      HitList hits = (*iclu)->predict(tre);
      // Fetch the cut record for the cluster.
      CutRecord* pcut_record2 = 0;
      if ( pcut_record1 != 0 ) {
        ClusterFilter::CutRecordMap::iterator icut = precs->find(*iclu);
        assert( icut != precs->end() );
        pcut_record2 = icut->second;
      }
      // Loop over hits
      // For each hit, copy track and refit with new hit.
      // Save new tracks if fit is successful.
      HitList::const_iterator ihit;
      for ( ihit=hits.begin(); ihit!=hits.end(); ++ihit ) {
        ++iaddhit_attempt;
        HitPtr phit = *ihit;
        // Report new hit.
        if ( prep != 0 ) {
          prep->make_report(*this, "PTrack add hit", phit);
        }
        // Construct a cut record for each hit.
        // If the fit is succesful, this will be assigned to the track.
	CutRecord* pcut_record3 = 0;
	if ( pcut_record2 != 0 ) {
	  pcut_record3 = pcut_record2->make_cut_record();
	}
        // refit the track with the new hit; save if fit succeeds
        MTrack trm = trp._trm;
        // If the hit and track have different surfaces, then propagate
        // the track to the hit surface.
        bool same_surface = true;
        const Hit& hit = **ihit;
        const Surface& srf_hit = hit.get_surface();
        const Surface& srf_trk = *trm.get_track().get_surface();
        string repstat = "unknown result";
        if ( srf_hit != srf_trk ) {
          if ( srf_hit.get_pure_type() == srf_trk.get_pure_type() ) {
            ETrack tre = trm.get_track();
            const Propagator& prop = cnl.get_propagator();
            PropStat pstat = trm.propagate(prop,srf_hit,Propagator::NEAREST);
            if ( ! pstat.success() ) {
              repstat = "propagation failed";
              same_surface = false;
            }
          } else {
            repstat = "track and hit had different surfaces";
            same_surface = false;
          }
        }
        if ( same_surface ) {
          bool fit_successful = 
            ! fit.add_hit_with_record(trm,phit,pcut_record3);
          repstat = "fit failed";
          if ( fit_successful )  {
            repstat = "check failed";
            PTrackPtr ptrp_new( 
              new PTrack(path,trm,chain,get_tuple(),pcut_record3) );
            if ( ptrp_new->check() ) {
              repstat = "successful prop, fit and check";
              ++iaddhit_success;
              new_tracks.push_back( ptrp_new );
              new_tracks.back()->set_parent_id( get_id() );
            }  // end successful fit check
          }  // end successful fit
        }  // end same surface check
        // Report propagation/fit status.
        if ( prep != 0 ) {
          prep->make_report(*this, "PTrack hit status", repstat);
        }
      }  // end loop over hits
    }  // end loop over clusters
    if ( fit_time != "" ) {
      int stat = ptimer->stop(fit_time);
      assert( stat == 0 );
    }

    // Add miss (if defined) to the original track.
    MissPtr pmiss = chain.get_miss();
    string repstat = "no miss defined";
    if ( pmiss ) {
      repstat = "miss added";
      ++iaddmiss_attempt;
      // If cut recording is turned on, copy the PTrack and do checking here
      // so that we get a unique cut record for the track with the added miss.
      if(!pcut_record1)
	trp._trm.add_miss(*pmiss);
      else {
	MTrack trm = trp._trm;
	trm.add_miss(*pmiss);
	CutRecord* pcut_record4 = pcut_record1->make_cut_record();
	PTrackPtr ptrp_miss(new PTrack(path,trm,chain,get_tuple(),
				       pcut_record4));
        repstat = "check failed";
	if ( ptrp_miss->check() ) {
	  ++iaddmiss_success;
	  new_tracks.push_back( ptrp_miss );
	  new_tracks.back()->set_parent_id( get_id() );
          repstat = "miss added and check successful";
	}
      }
    }

    // We keep the original track (i.e. that with no clusters added) if
    // 1. a miss has been added and cut recording is off or
    // 2. there were no clusters associated with the layer surface
    if ( original_time != "" ) {
      int stat = ptimer->start(original_time);
      assert( stat == 0 );
    }
    bool keep_original = 
      pmiss != 0 && pcut_record1 == 0 || pmiss==0 && !chain.has_clusters();

    // Check the original track and delete if check fails.
    // Note that track with no miss now survives.
    if ( keep_original && trp.check() ) {
      ++itrp;
      ++iaddmiss_success;
    } else {
      new_tracks.erase(itrp++);
      repstat += "--track deleted";
    }
    // Report original track status.
    if ( prep != 0 ) {
      prep->make_report(*this, "PTrack original track status", repstat);
    }
    if ( original_time != "" ) {
      int stat = ptimer->stop(original_time);
      assert( stat == 0 );
    }

    if(precs)
      delete precs;

  }  // end loop over new tracks

  // Loop over new tracks and refit and apply post checkers.
  if ( post_time != "" ) {
    int stat = ptimer->start(post_time);
    assert( stat == 0 );
  }
  itrp = new_tracks.begin();
  while ( itrp != new_tracks.end() ) {
    string repstat = "no refit or post check";
    PTrack& trp = **itrp;
    const Path& path = trp.get_path();
    // Create pointer to the PTrack. This will be assigned a new value if
    // the track is refit and set null if the track is to be dropped.
    PTrackPtr& ptrp = *itrp;
    bool keep = true;
    // Refit.
    // Report refit track.
    if ( prep != 0 ) {
      prep->make_report(*this, "PTrack refit", *ptrp);
    }
    if ( path.has_refitter() ) {
      MTrack& trm = trp.get_mutable_track();
      int stat = path.get_refitter().fit(trm);
      repstat = "refit ok";
      if ( stat != 0 ) {
        repstat = "refit failed";
        keep = false;
      }
    }
    // Apply post checkers.
    if ( keep && !ptrp->post_check() ) {
      keep = false;
      repstat = "post check failed";
    }
    if ( keep ) {
      ++itrp;
    } else {
      new_tracks.erase(itrp++);
      repstat += "--track deleted";
    }
    // Report refit and post check status.
    if ( prep != 0 ) {
      prep->make_report(*this, "PTrack refit status", repstat);
    }
  }
  if ( post_time != "" ) {
    int stat = ptimer->stop(post_time);
    assert( stat == 0 );
  }
        
  // Loop over new tracks and set stopped if the path is
  // 1. at the exit of the layer and
  // 2. is at a stop or has no children.
  for ( itrp=new_tracks.begin(); itrp!=new_tracks.end(); ++itrp ) {
    PTrack& trp = **itrp;
    if ( trp._chain.get_current_status().at_exit() &&
         ( trp.get_path().get_stop() || 
           trp.get_path().get_children().size() == 0 )  )
      trp._stopped = true;
  }  // end loop over new tracks

  // Fill the ntuple.
  if ( tuple_is_active() ) {
    _ptup->capture( "evid", EventID::get_global_event_number() );
    _ptup->capture( "path", _path.get_id() );       // path ID
    _ptup->capture( "npath", ichildpath );          // # child paths
    _ptup->capture( "nprop", iprop );               // # layer prop's
    _ptup->capture( "nfittry", iaddhit_attempt );   // # fit attempts
    _ptup->capture( "nfitok", iaddhit_success );    // # fit successes
    _ptup->capture( "nmisstry", iaddmiss_attempt ); // # miss successes
    _ptup->capture( "nmissok", iaddmiss_success );  // # miss successes
    TrackMcHitInfo trkmcinfo = get_track().get_mcinfo();
    _ptup->capture( "mcid", trkmcinfo.best_mcid() );  // best MC id
    _ptup->capture( "mcpurity", float(trkmcinfo.purity()) );  // purity of that id
    const TrackVector& vec = get_track().get_track().get_vector();
    _ptup->capture( "t0", float(vec(0)) );
    _ptup->capture( "t1", float(vec(1)) );
    _ptup->capture( "t2", float(vec(2)) );
    _ptup->capture( "t3", float(vec(3)) );
    _ptup->capture( "t4", float(vec(4)) );
    _ptup->storeCapturedData();
    _ptup->clearData();
  }

  // Write debugging information.
  if ( debug ) {
    cout << "# child paths: " << ichildpath << new_line;
    cout << "# propagations: " << iprop << new_line;
    cout << "# fit successes/attempts: "
         << iaddhit_success << " / " << iaddhit_attempt << new_line;
    cout << "# miss successes/attempts: "
         << iaddmiss_success << " / " << iaddmiss_attempt << new_line;
  }

  // Stop the total clock.
  if ( total_time != "" ) {
    int stat = ptimer->stop(total_time);
    assert( stat == 0 );
  }

  // Return the new tracks.
  return new_tracks;
}

//************************************************************************

// clear the stopped flag

void PTrack::clear_stop() {
  _stopped = false;
}

//************************************************************************

// Set and return the end of path flag.

bool PTrack::check_end_of_path() {

  // If the current layer status is defined, then
  // track must be at exit of its layer.
  if ( _chain.cluster_status_is_valid() &&
       ! _chain.get_current_status().at_exit() )
    return _end_of_path = false;

  // Stops must be cleared.
  if ( is_stopped() ) return _end_of_path = false;

  // Track must have no children or be at the end of a path.
  if (  _path.get_children().size() )
    if (  ! _path.get_end()  ) return _end_of_path = false;

  return _end_of_path = true;

}
    
//************************************************************************

// Create an ntuple.

int PTrack::set_tuple(const char* name) {
  // If the name is blank, exit.
  if ( string(name).length() == 0  ) return 1;
  // If global tuple flag is not enabled, exit.
  if ( ! tuple_is_enabled() ) return 2;
  // create the ntuple
  _ptup = &TupleManager::tuple(name);
  // define the columns
  float xdefault = -9999.99;
  _ptup->columnDirect( "evid", -1 );
  _ptup->columnDirect( "path", -1 );
  _ptup->columnDirect( "npath", -1 );
  _ptup->columnDirect( "nprop", -1 );
  _ptup->columnDirect( "nfittry", -1 );
  _ptup->columnDirect( "nfitok", -1 );
  _ptup->columnDirect( "nmisstry", -1 );
  _ptup->columnDirect( "nmissok", -1 );
  _ptup->columnDirect( "mcid", -1 );
  _ptup->columnDirect( "mcpurity", xdefault );
  _ptup->columnDirect( "t0", xdefault );
  _ptup->columnDirect( "t1", xdefault );
  _ptup->columnDirect( "t2", xdefault );
  _ptup->columnDirect( "t3", xdefault );
  _ptup->columnDirect( "t4", xdefault );

  return 0;
}

//************************************************************************

// Set the parent id.
void PTrack::set_parent_id(int id) {
  _parent_id = id;
}

//************************************************************************
// friends
//************************************************************************

// output stream

ostream& trf::operator<<(ostream& stream, PTrack& trp) {
  trp.ostr(stream);
  return stream;
}

//************************************************************************