// findcyl.cpp

// Main program to find tracks in a cylindrical detector.

#include <strstream>
#include <iostream>
#include <list>
#include "trfutil/trfstream.h"
#include "ptr/Ptr.h"
#include "D0Detector.h"
#include "D0DetectorSimulator.h"
#include "MCTrackMaker.h"
#include "trf_cft/CFTPathBuilder.h"
#include "trfcyl/SurfCylinder.h"
#include "trfbase/ETrack.h"
#include "trffind/TrackFinder.h"
#include "trfcyl/PropCyl.h"
#include "trfbase/PropStat.h"
#include "trfcyl/CylTrackNtuple.h"
#include "trffind/PTrack.h"
#include "trffit/RTrack.h"
#include "trfbase/MCTrack.h"
#include "trffit/TrackMatch.h"
#include "trffit/FullFitKalman.h"
#include "trfutil/RandomRegistry.h"
#include "trfutil/EventID.h"
#include "trfcyl/ClusterFindCyl.h"
#include "trfcyl/AddFitCyl_Phi_Phi_Phi.h"
#include "trfcyl/AddFitCyl_Phi_Phi.h"
#include "trfcyl/AddFitCyl_PhiZ.h"
#include "trfbase/VTrackGenerator.h"
#include "trffind/PathTree.h"
#include "trffit/register_trffit_types.h"
#include "trffind/register_trffind_types.h"
#include "trfcyl/register_trfcyl_types.h"

using std::istrstream;
using std::cin;
using std::cout;
using std::cerr;
using std::endl;
using std::list;

using namespace trf;

// Assign track parameter indices.
enum {
  IPHI = SurfCylinder::IPHI,
  IZ   = SurfCylinder::IZ,
  IALF = SurfCylinder::IALF,
  ITLM = SurfCylinder::ITLM,
  IQPT  = SurfCylinder::IQPT
};

int main(int argc, char* argv[]) {

  cout << trf_format;

  //********************************************************************
  // Handle input arguments.  Evaluate nevt, ntrk and nprint.
  //********************************************************************

  int nevt = 10;
  int ntrk = 1;
  int nprint = 0;
  if ( argc > 1 ) {
    char* cevt = argv[1];
    istrstream(cevt,0) >> nevt;
  }
  if ( argc > 2 ) {
    char* ctrack = argv[2];
    istrstream(ctrack,0) >> ntrk;
  }
  if ( argc > 3 ) {
    char* cprint = argv[3];
    istrstream(cprint,0) >> nprint;
  }

  // Initialize the print flag.
  // For nprint < 0, nothing is printed.
  bool print = nprint >=0;

  //********************************************************************
  // Register types with ObjTable.
  //********************************************************************

  register_trffit_types();
  register_trffind_types();
  register_trfcyl_types();

  //********************************************************************
  // Construct input track generator.
  //********************************************************************

  MCTrackMaker track_maker;

  if ( print ) {
    cout << endl;
    cout << "Monte Carlo track generator:" << endl;
    cout << track_maker << endl;
  }

  //********************************************************************
  // Set the parameters.
  //********************************************************************

  const double bfield = 2.0;

  // For 2-phi fitter.
  const double ptmin = 0.45;

  // max chi-square diff for cluster finders
  //const double max_chsq_diff = 20.0;

  // maximum sqrt(chi-square diff) for 3-phi fitter.
  const double nsigma_3phi = 5.0;

  // Enable ntuples.
  //ClusterFindCyl::enable_tuple();
  //AddFitCyl_Phi_Phi::enable_tuple();
  //AddFitCyl_Phi_Phi_Phi::enable_tuple();
  //AddFitCyl_PhiZ::enable_tuple();
  //PTrack::enable_tuple();

  //********************************************************************
  // Construct layer manager.
  //********************************************************************

  DetectorPtr pdet;
  pdet = new D0Detector;
  const Detector& det = *pdet;
  D0DetectorSimulator detsim(pdet);

  if ( print ) cout << endl << det << endl;
  if ( print ) cout << endl << detsim << endl;

  //********************************************************************
  // Define the starting track.
  //********************************************************************

  // Define the starting surface for the fit tracks.
  double r0 = 60.0;
  SurfacePtr psrf0( new SurfCylinder(r0) );

  // Define the starting track vector.
  TrackVector vec0;

  // Define the starting error matrix.
  // This must be appropriate for the above surface.
  TrackError err0;
  err0(IPHI,IPHI) =  20.0;
  err0(IZ,IZ)     = 100.0;
  err0(IALF,IALF) =  20.1;
  err0(ITLM,ITLM) =  10.0;
  err0(IQPT,IQPT) =  20.0;

  // Use the above to define the starting ETrack.
  ETrack tre0(psrf0,vec0,err0);

  //********************************************************************
  // Build paths.
  //********************************************************************

  CFTPathBuilder path_maker( det, tre0, bfield, ptmin,
                             nsigma_3phi );
  MutablePathPtr ppath0( new Path );
  path_maker.add_path_tree(ppath0);
  const Path& path0 = *ppath0;
  if (print ) cout << endl << PathTree(ppath0) << endl;

  //********************************************************************
  // Build the track finder.
  //********************************************************************

  // Build the track finder.
  TrackFinder tfind(path0,tre0);
  assert( path_maker.get_stops().find("end") != 
          path_maker.get_stops().end() );
  tfind.add_stop( *(*path_maker.get_stops().find("end")).second );
  tfind.set_tuple("ptrack");
  //tfind.set_debug(1);
  if ( print ) cout << endl << tfind << endl;

  //********************************************************************
  // Build fitter for refitting tracks.
  //********************************************************************

  // Define propagator.
  PropagatorPtr pprop(new PropCyl(ObjDoublePtr(new ObjDouble(bfield))));

  // Define fitter.
  FullFitKalman fullfit(pprop);

  // Should we refit?
  bool refit = false;

  //********************************************************************
  // Build the propagator and surface for matching tracks.
  //********************************************************************

  // Define surface.
  // We use a pointer so we can override this definition later.
  SurfacePtr psrfm( new SurfCylinder(50.0) );

  //********************************************************************
  // Create the ntuple.
  //********************************************************************

  CylTrackNtuple ntp(true,true,false);

  // Construct arrays to hold the random number generators.

  //********************************************************************
  // Construct random number registry.
  //********************************************************************

  // This registry record the state of the random number generators for
  // each event and can be used to recover this state for any event.
  // This simplifies debugging.

  RandomRegistry random_state;
  
  // Register the track maker.
  track_maker.register_generators(random_state);
  assert( random_state.get_generator_count() == 1 );

  // Register the hit generators.
  detsim.register_generators(random_state);
  assert( random_state.get_generator_count() == 17 );

  // Record the inital state.
  int reg_flag = random_state.record();
  assert( reg_flag == 0 );

  //********************************************************************
  // Set the interactive status.
  //********************************************************************

  // If nevt <=0, then run in interactive mode:
  // After processing -nevt events, prompt user for the next event.
  // If this event is beyond the range, continue to the new event.
  // If the event has already been processed, reprocess it only.
  bool interactive = false;
  if ( nevt <= 0 ) {
    interactive = true;
    nevt = -nevt;
  }

  // Intitialize the flag indicating that only one event should be
  // processed.
  bool single_event = false;

  //********************************************************************
  // Loop over events
  //********************************************************************

  // initialize current event number
  int ievt = 0;

  // intialize last event number
  int lastevt = 0;

  while ( true ) {

    // Increment the event number
    ievt = lastevt + 1;

    // and take special action if all events have been processed or
    // if we are in single event mode.
    if ( ievt > nevt  ||  single_event )  {

      // If we are in interactive mode, fetch the event number.
      if ( interactive ) {
        cout << "Event number (0 to exit)> ";
        int newevt;
        cin >> newevt;
        // event < 0 ==> exit
        if ( newevt <= 0 ) break;
        // event > nevt ==> reset nevt and continue to newevt
        if ( newevt > nevt ) {
          ievt = nevt + 1;
          nevt = newevt;
          single_event = false;
        }
        // event already processed ==> reprocess
        else {
          ievt = newevt;
          single_event = true;
        }
      }

      // If not in interactive mode, exit.
      else break;

    }

    // If the last event processed was not ievt-1, then we need to
    // reset the random generators to the latter state.
    if ( ievt != lastevt+1 ) {
      assert( random_state.get_state_count() >= ievt );
      random_state.set(ievt-1);
    } else {
      assert( random_state.get_state_count() == ievt ||
              random_state.get_state_count() == nevt+1  );
    }

    assert( ievt <= nevt );

    // Evaluate whether this event should be displayed.
    print = ievt<=nprint || single_event;

    if ( print ) cout << endl << "Begin event " << ievt << endl;

    // Build event identifier.
    // The event number may be obtained with
    // EventID::get_global_event_number().
    EventID evid(ievt);
    assert( EventID::get_global_event_number() == ievt );

    //******************************************************************
    // Generate tracks.
    //******************************************************************

    MCTrackList mctracks;
    int i;
    for ( i=0; i<ntrk; ++i ) 
      mctracks.push_back( MCTrackPtr( track_maker.new_track() ) );

    if ( print ) {
      cout << "**************** MC Tracks *****************" << endl;
      MCTrackList::const_iterator itrv;
      for ( itrv=mctracks.begin(); itrv!=mctracks.end(); ++itrv )
        cout << **itrv << endl;
    }

    //******************************************************************
    // Generate clusters and store them on the layers.
    //******************************************************************

    // Drop the old clusters.
    detsim.drop_clusters();

    // Loop over MC tracks and add new clusters for each.
    MCTrackList::const_iterator itrmc;
    for ( itrmc=mctracks.begin(); itrmc!=mctracks.end(); ++itrmc )
      detsim.add_clusters(**itrmc);

    if ( print ) {
      cout << "************ Layers with hits **************" << endl;
      det.print_layers();
      cout << endl;
    }

    //******************************************************************
    // Find tracks.
    //******************************************************************

    // Find.
    PTrack::PTrackList ptracks = tfind.find();

    // Extract the reconstructed tracks (MTrack's).
    RTrackList rtracks;
    PTrack::PTrackIterator itrp;
    for ( itrp=ptracks.begin(); itrp!=ptracks.end(); ++itrp )
      rtracks.push_back( RTrackPtr( new RTrack( (*itrp)->get_track() ) ) );

    if ( print ) {
      int ntrack = ptracks.size();
      cout << "************ " << ntrack
           <<                " track";
      if ( ntrack != 1 ) cout << "s";
      cout <<                       " found ***************" << endl;
      string line = 
        "-------------------------------------------------------";
      cout << line << endl;
      PTrack::PTrackIterator itrp;
      for ( itrp=ptracks.begin();  itrp!=ptracks.end();  ++itrp )
        cout << **itrp << endl << line << endl;
    }

    //******************************************************************
    // Refit tracks.
    //******************************************************************

    if ( refit ) {

      if ( print ) {
        cout << "****************** Refit tracks *******************";
        cout << endl;
      }

      RTrackList::const_iterator itrr;
      for ( itrr=rtracks.begin(); itrr!=rtracks.end(); ++itrr ) {
        HTrack& trh = **itrr;
        // Calculate starting track parameters by propagating fit
        // track back to the starting surface.
        ETrack tre = trh.get_track();
        PropStat pstat = pprop->vec_prop(tre,*psrf0);
        // If propagation fails, keep original track.
        if ( ! pstat.success() ) {
          cerr << "Refit propagation failed." << endl;
          continue;
        }
        // Use the standard starting errors.
        tre.set_error( tre0.get_error() );
        // Do a full fit with the new starting parameters.
        trh.set_fit(tre,0.0);
        fullfit.fit(trh);
        if ( print ) {
          cout << trh << endl;
          cout <<
            "-------------------------------------------------------";
          cout << endl;
        }
      }  // end loop over tracks

    }  // end refit

    //******************************************************************
    // Propagate tracks to a common surface.
    //******************************************************************

    // Redefine the match surface to be that of the first reconstructed
    // track.
    if ( rtracks.size() ) psrfm = 
      rtracks.front()->get_track().get_surface();

    // If propagation fails, we leave track at its original surface.
    // The track match quality will be invalid if tracks are at different
    // surfaces.

    // Propagate MC tracks.
    MCTrackList::iterator itrm;
    for ( itrm=mctracks.begin(); itrm!=mctracks.end(); ++itrm ) {
      PropStat pstat = pprop->vec_prop( **itrm, *psrfm );
      if ( ! pstat.success() ) {
        if ( print ) 
          cout << "***** MC track match propagation failed!!!!" << endl;
      }
    }

    // Propagate reconstructed tracks.
    RTrackList::iterator itrr;
    for ( itrr=rtracks.begin(); itrr!=rtracks.end(); ++itrr ) {
      PropStat pstat = (*itrr)->propagate(*pprop,*psrfm);
      if ( ! pstat.success() ) {
        if ( print ) 
          cout << "***** Reco track match propagation failed!!!!"
               << endl;
      }
    }

    //******************************************************************
    // Match tracks.
    //******************************************************************

    TrackMatch match(mctracks,rtracks);

    if ( print ) {
      cout << "************* Match tracks **************" << endl;
      cout << match << endl;
    }

    //******************************************************************
    // Fill ntuple.
    //******************************************************************

    ntp.set_event(ievt);
    ntp.fill_match(match);

    //******************************************************************
    // Record random numbers for the next event.
    //******************************************************************

    // If this is the first time this event has been processed,
    // record the random number state.
    // Random registry index i holds the state after event i, i.e
    // the state before event i+1.
    // Registry holds the initial state plus one state for each event.
    if ( ievt >= random_state.get_state_count() ) {
      reg_flag = random_state.record();
      assert( reg_flag == ievt );
    }

    //******************************************************************
    // End event.
    //******************************************************************

    lastevt = ievt;
    if ( print ) cout << endl << "End of event " << ievt << endl;

  }  // end loop over events

}