// TrackExtender.cpp

#include "TrackExtender.h"
#include <iostream>
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfobj/Message.h"
#include "trfutil/trfstream.h"
#include "trfutil/CpuTimerTuple.h"
#include "trfcut/CutRecorder.hpp"
#include "Path.h"
#include "PTrack.h"
#include "PathStop.h"
#include "Filter.h"

using std::cout;
using std::cerr;
using std::endl;

using trf::Message;
using trf::PathStop;
using trf::Filter;
using trf::CutRecord;
using trf::CutRecorder;
using trf::PTrack;
using trf::TrackExtender;

//**********************************************************************
// typedefs
//**********************************************************************

// Array for PTrack pointers.
typedef vector<PTrack::PTrackPtr> PTrackArray;

//**********************************************************************
// Free functions.
//**********************************************************************

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "TrackExtender" );
  TrackExtender::PathStopList stops;
#ifdef ObjData_supports_lists
  data.get_bare_ptr_list("stops",stops);
#endif
  TrackExtender* pxtn = new TrackExtender;
  for ( TrackExtender::PathStopList::const_iterator istp=stops.begin();
        istp!=stops.end(); ++istp ) pxtn->add_stop(**istp);
  string tuple_name = data.get_string( "tuple_name" );
  pxtn->set_tuple(tuple_name);
  return ObjPtr(pxtn);
}

}

//**********************************************************************
// methods
//**********************************************************************

// output stream

void TrackExtender::ostr(ostream& stream) const {
  stream << begin_object;
  stream << "Global track extender has ";
  int size = _stops.size();
  stream << size << " stop";
  if ( size != 1 ) stream << "s";
  if ( size > 0 ) stream << ":";
  else stream << ".";
  for ( PathStopConstIterator istop=_stops.begin();
        istop!=_stops.end(); ++istop ) stream << new_line << **istop;
  if ( _tuple_name.length() ) {
    stream << "Ntuple name: " << _tuple_name;
  } else {
    stream << "No ntuple.";
  }
  stream << end_object;
}

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

// constructor

TrackExtender::TrackExtender(CpuTimerTuple* ptimer)
: _ptimer(ptimer), _tuple_name("") { }

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

// destructor

TrackExtender::~TrackExtender() { }

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

// Return the creator.

ObjCreator TrackExtender::get_creator() {
  return create;
}

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

// Write the object data.

ObjData TrackExtender::write_data() const {
  ObjData data( get_type_name() );
#ifdef ObjData_supports_lists
  data.add_bare_ptr_list( "stops", get_stops() );
#endif
  data.add_string( "tuple_name", get_tuple_name() );
  return data;
}

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

// add a stop

void TrackExtender::add_stop(const PathStop& stop) {
  _stops.push_back(PathStopPtr(&stop));
}

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

// Define ntuple.  If this is called with a nonzero length argument
// then A PTrack ntuple with the given name will be filled.

void TrackExtender::set_tuple(string name) {
  _tuple_name = name;
}

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

// Extend a list of tracks.

void TrackExtender::extend(PTrack::PTrackList& tracks) const {

  if ( get_debug() ) {
    cout << "-----------------------------------------------------" << endl;
    cout << "TrackExtender called with " << tracks.size() << " tracks." << endl;
    cout << "-----------------------------------------------------" << endl;
  }

  // Report the initial tracks.
  make_report(*this, "begin", tracks);

  // Set the tuple for all input tracks.
  if ( _tuple_name.length() ) {
    for ( PTrack::PTrackIterator itrp=tracks.begin();
          itrp!=tracks.end(); ++itrp ) {
      (**itrp).set_tuple(_tuple_name.c_str());
    }
  }

  // Initialize track iterator.
  PTrack::PTrackIterator itrp = tracks.begin();

  // Loop over stops
  // Loop once more after all stops are processed.
  PathStopConstIterator istop = _stops.begin();
  int nloop = 0;
  int nstop = 0;
  while ( true ) {
    ++nloop;

    // Report the initial tracks  and the stop for this stop.
    make_report(*this, "begin stop propagation stop", **istop);
    make_report(*this, "begin stop propagation", tracks);

    // Loop over tracks.
    while ( itrp != tracks.end() ) {
      PTrack& trp = **itrp;

      if ( get_debug() ) {
        cout << "-----------------------------------------------------" << endl;
        cout << "Old track:" << endl;
        cout << "-----------------------------------------------------" << endl;
        cout << trp << endl;
        cout << "-----------------------------------------------------" << endl;
      }

      // Propagate the track to produce new tracks.
      // If track is stopped or at its end of path, go to the next.
      if ( trp.is_stopped() || trp.at_end_of_path() ) {
        ++itrp;
        ++nstop;
        if ( get_debug() ) {
          cout <<
            "-----------------------------------------------------" << endl;
          if ( trp.is_stopped() ) cout << "Skipping track at stop." << endl;
          else cout << "Skipping track at end of path." << endl;
          cout <<
            "-----------------------------------------------------" << endl;
        }
        continue;
      }

      // Report the track to be propagated.
      make_report(*this, "next", trp);

      const TrfReportingObject* prep = 0;
      if ( reporters().size() ) prep = this;
      PTrack::PTrackList new_tracks = trp.propagate(get_debug(), _ptimer, prep);

      // Report the new tracks to be added to the list.
      make_report(*this, "new tracks", new_tracks);

      // Add the new tracks to the list.
      // We put them after the current track so they will be
      // processed immediately.
      {
        PTrack::PTrackList::iterator itrp_insert = itrp;
        ++itrp_insert;
        tracks.insert( itrp_insert, new_tracks.begin(), 
                                    new_tracks.end() );
      }

      // Drop the old track and set iterator to point at the next.
      // If the track is at an end of path, leave it on the list.
      if ( trp.check_end_of_path() ) {
        ++itrp;
        ++nstop;
      } else {
        tracks.erase(itrp++);
      }

      // Report the new track list.
      make_report(*this, "after prop", tracks);

      if ( get_debug() ) {
        cout << "-----------------------------------------------------" << endl;
        cout << new_tracks.size() << " new tracks found." << endl;
        cout << tracks.size() << " remain on list." << endl;
        cout << nstop << " tracks are stopped." << endl;
        cout << "-----------------------------------------------------" << endl;
      }

      // If the track list gets too large, give up.
      if ( tracks.size() > 100000 ) {
        make_report(*this, "too many tracks", tracks);
        tracks.erase( tracks.begin(), tracks.end() );
        cout << "TrackExtender: Too many tracks--giving up." << endl;
        cerr << "TrackExtender: Too many tracks--giving up." << endl;
        assert( tracks.size() == 0 );
        return;
      }

    }  // end loop over tracks

    // exit loop if there are no more stops.
    if ( istop == _stops.end() ) {
      break;
    }

    // Report the final tracks for this stop.
    make_report(*this, "end stop propagation", tracks);

    if ( get_debug() ) {
      cout << "-----------------------------------------------------" << endl;
      cout << "Processing stop " << nloop << "." << endl;
      assert( nstop == tracks.size() );
    }

    // create reference to the current stop
    const PathStop& stop = **istop;

    // Loop over tracks.
    // Any at the current stop are transferred to the stopped list.
    PTrack::PTrackList stop_tracks;
    itrp = tracks.begin();
    while ( itrp != tracks.end() ) {
      PTrack& trp = **itrp;
      if ( trp.get_path().get_stop() == &stop ) {
        stop_tracks.push_back(*itrp);
        tracks.erase(itrp++);
      }
      else ++itrp;
    }  // end loop over tracks
    nstop = tracks.size();
    int nfilter = stop_tracks.size();

    // Report the final tracks for this stop.
    make_report(*this, "begin filter", stop_tracks);

    // Filter the stopped tracks.
    Filter::FlagArray flags;
    const PathStop::FilterList& filters = stop.get_filters();
    PathStop::FilterList::const_iterator ifil;
    // Loop over filters.
    int npass = 0;
    int nfail = 0;
    for ( ifil=filters.begin(); ifil!=filters.end(); ++ifil ) {
      const Filter& filt = **ifil;
      // Build list of miss tracks.
      Filter::MTrackArray stop_mtracks;
      Filter::CutRecordArray* precs=0;
      for (itrp=stop_tracks.begin(); itrp!=stop_tracks.end(); ++itrp) {
        PTrack& trp = **itrp;
        stop_mtracks.push_back( Filter::MTrackPtr(&trp.get_track()) );
	CutRecord* prec = trp.get_cut_record();
	if(prec) {
	  if(!precs) {
	    precs = new Filter::CutRecordArray;
	    precs->reserve(stop_mtracks.size());
	  }
	  precs->push_back(prec);
	}
      }
      // apply the filter
      assert(precs == 0 || precs->size() == stop_mtracks.size());
      Filter::FlagArray flags = filt.process_with_records(stop_mtracks, precs);
      // Loop over tracks. Refit and apply checkers. Drop tracks that fail the
      // filter, the fit or any of the checkers.
      itrp = stop_tracks.begin();
      for ( int i=0; i<flags.size(); ++i ) {
        PTrack::PTrackPtr& ptrp = *itrp;
        make_report(*this, "filter", **itrp);
        bool keep = flags[i];
        string repstat = "filter failure--track deleted";
        if ( keep ) {
          MTrack& trm = ptrp->get_mutable_track();
          int stat = stop.fitcheck(trm);
          if ( stat != 0 ) {
            repstat = "fit/check failure--track deleted";
            keep = false;
          }
        }
        if ( keep ) {
          repstat = "track ok";
          ++itrp;
        } else {
          stop_tracks.erase(itrp++);
        }
        make_report(*this, "filter status", repstat);
      }
      assert( itrp == stop_tracks.end() );
      if ( precs ) {
	delete precs;
      }
    }  // end loop over filters
    int nkeep = stop_tracks.size();

    // Clear the stops.
    PTrack::PTrackIterator jtrp;
    for ( jtrp=stop_tracks.begin(); jtrp!=stop_tracks.end(); ++jtrp )
      (*jtrp)->clear_stop();
    
    // Report the final tracks for this stop.
    make_report(*this, "end filter", stop_tracks);

    // Tracks that survive the filters are copied back to the end
    // of the track list.
    // Set itrp to the first of the copied tracks.
    bool has_old_tracks = tracks.size();
    if ( has_old_tracks ) {
      itrp = tracks.end();
      --itrp;
    }
    tracks.splice( tracks.end(), stop_tracks );
    if ( has_old_tracks ) ++itrp;
    else itrp = tracks.begin();

    ++istop;

    if ( get_debug() ) {
      cout << nkeep << " of " << nfilter << " tracks survive." << endl;
      cout << tracks.size() << " tracks remain on list." << endl;
      cout << "-----------------------------------------------------" << endl;
    }

  }  // end loop over stops

  if ( get_debug() ) {
    cout << "-----------------------------------------------------" << endl;
    cout << "All stops processed." << endl;
    cout << tracks.size() << " tracks remain on list." << endl;
    cout << "-----------------------------------------------------" << endl;
  }

  // Report the final tracks.
  make_report(*this, "end", tracks);

  // If cut recording is enabled, set the good track flag for the cut records
  // of any remaining tracks.

  CutRecorder* prec = CutRecorder::instance();
  if(prec->enabled()) {
    for(PTrack::PTrackList::const_iterator i = tracks.begin();
	i != tracks.end();
	++i) {
      CutRecord* p = (*i)->get_cut_record();
      if(p)
	p->set_good_track();
    }
  }
}
  
//**********************************************************************