// fitcyl.cpp

// Main program to fit tracks with cylinder hits.

#include "CLHEP/Random/RandFlat.h"
#include <strstream>
#include "trfutil/trfstream.h"
#include "trffit/HTrackGenerator.h"
#include "trfcyl/HitCylPhiGenerator.h"
#include "trfcyl/HitCylPhiZGenerator.h"
#include "trfcyl/PropCyl.h"
#include "trfcyl/SurfCylinder.h"
#include "trfbase/TrackVector.h"
#include "trfbase/VTrackGenerator.h"
#include "trfutil/TRFMath.h"
#include "ptr/Ptr.h"
#include "ptr/DeletePolicy.h"
#include "trfbase/VTrack.h"
#include "trffit/HTrack.h"
#include "trffit/FullFitKalman.h"
#include "trfbase/PropStat.h"
#include "trfcyl/CylTrackNtuple.h"
#include "trfutil/RandomRegistry.h"

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

using namespace trf;

// Random number seeds.
long myseed(int i) {
  return RandFlat::shootInt(100000000);
}

// 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 ntrack and nprint.
  //********************************************************************

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

  // cerr << "Usage: " << argv[0] << " ntrack nprint" << endl;
  // return 1;

  //********************************************************************
  // Construct hit track generator.
  //********************************************************************

  // Define the hit generators.
  // Geometry is D0 scifi (Aug 97)
  double dx = 0.01;  // 0.01 cm = 100 um
  HTrackGenerator::HitGeneratorList hgens;
  // define radii for acial measurements
  int nbarrel = 8;
  double rcyl[]={ 19.50, 23.41, 28.09, 32.77, 37.46, 42.14, 48.75, 51.50};
  // phi msmts
  int ib;
  for ( ib=0; ib<nbarrel; ++ib ) {
    double r = rcyl[ib];
    hgens.push_back( HTrackGenerator::HitGeneratorPtr(
                     new HitCylPhiGenerator( r, dx/r, myseed(ib) ) ) );
  }
  // stereo msmts
  double stereo = 0.001;
  for ( ib=0; ib<nbarrel; ++ib ) {
    stereo *= -1.0;
    double r = rcyl[ib] + 0.2;
    hgens.push_back( HTrackGenerator::HitGeneratorPtr(
             new HitCylPhiZGenerator( r, stereo, dx/r, myseed(8+ib) ) ) );
  }

  // Define the propagator
  PropagatorPtr pprop( new PropCyl(ObjDoublePtr(new ObjDouble(2.0))));

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

  // 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) =  10.0;

  // Construct the hit track generator.
  HTrackGenerator htgen(hgens,*pprop,srf0,err0);

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

  // Define the starting surface.
  double ri = 1.0;
  SurfCylinder srfi( ri );

  // Define lower track bounds.
  TrackVector tmin;
  tmin(IPHI) = -PI;       // phi
  tmin(IZ)   = -30.0;     // z
  tmin(IALF) = -PI2;      // alpha = phi_dir - phi_pos
  tmin(ITLM) = -1.0;      // tan(lambda)
  tmin(IQPT) = -0.001;    // q/p in e/GeV/c

  // Define upper track bound.
  TrackVector tmax;
  for ( int i=0; i<5; ++i ) tmax(i) = -tmin(i);

  // Construct the input track generator.
  VTrackGenerator vtgen( srfi, tmin, tmax );

  //********************************************************************
  // 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 generator.
  random_state.add_generator(vtgen);
  assert( random_state.get_generator_count() == 1 );

  // Register the hit generators.
  HTrackGenerator::HitGeneratorList::const_iterator igen;
  for ( igen=hgens.begin(); igen!=hgens.end(); ++igen )
    random_state.add_generator( **igen );
  assert( random_state.get_generator_count() == 17 );

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

  //********************************************************************
  // Construct the track fitter.
  //********************************************************************

  FullFitKalman fitk(pprop);

  //********************************************************************
  // Open HBOOK file and define track ntuple.
  //********************************************************************

  CylTrackNtuple ntp(true,true,false);

  //********************************************************************
  // Define the surface at which tracks are compared.
  //********************************************************************

  SurfCylinder srfc(20.0);

  //********************************************************************
  // Begin loop over tracks.
  //********************************************************************

  // First process |ntrack| events and then if ntrack <= 0, go into
  // interactive mode.

  Ptr<VTrack,DeletePolicy> ptrv;
  Ptr<HTrack,DeletePolicy> ptrh;
  PropStat pstat;
  int itrack = 0;
  bool interactive = false;
  bool check_interactive = false;
  if ( ntrack <= 0 ) {
    ntrack = -ntrack;
    check_interactive = true;
  }
  while ( itrack < ntrack  ||  check_interactive ) {
    ++itrack;
    // If we reach the end of the list, go into interactive mode.
    if ( itrack > ntrack ) interactive = true;

    // If in interactive mode, interrogate user for the next event.
    if ( interactive ) {
      cout << "Event number (0 to exit)> ";
      int newevt;
      cin >> newevt;
      // event < 0 ==> exit
      if ( newevt <= 0 ) break;
      // event > ntrack ==> reset ntrack and continue to newevt
      if ( newevt > ntrack ) {
        itrack = ntrack + 1;
        ntrack = newevt;
        interactive = false;
      }
      // event already processed ==> reprocess
      else {
        itrack = newevt;
      }
    }

    // If in interactive mode reset random state.
    if ( interactive ) {
      assert( random_state.get_state_count() >= itrack );
      random_state.set(itrack-1);
    }

    // set print flag
    bool lprint = itrack < nprint  ||  interactive;

    if ( lprint ) cout << endl;
    // generate an input track
    ptrv = vtgen.new_track();
    if ( lprint ) cout << *ptrv << endl;
    // generate a track with random hits
    ptrh = htgen.new_track(*ptrv);
    if ( lprint ) cout << *ptrh << endl;
    // Set starting values for track fit
    ETrack tre = ptrh->get_track();
    TrackVector vec = tre.get_vector();
    //vec(0) = 0.0;
    //vec(1) = 0.0;
    //vec(2) = 0.0;
    //vec(3) = 0.0;
    //vec(4) = 0.0;
    tre.set_vector(vec);
    ptrh->unset_fit(tre);
    if ( lprint ) cout << *ptrh << endl;
    // fit the track
    fitk.fit(*ptrh);
    // propagate the starting track to the comparison surface
    pstat = pprop->vec_prop(*ptrv,srfc);
    if ( ! pstat.success() ) {
      cerr << "Starting propagation failed."  << endl;
      continue;
    }
    // propagate the fitted track to the input surface
    pstat = ptrh->propagate(*pprop,srfc);
    if ( ! pstat.success() ) {
      cerr << "Fit propagation failed."  << endl;
      continue;
    }
    if ( lprint ) cout << *ptrh << endl;
    // fill ntuple
    ntp.set_event(itrack);
    ntp.fill(*ptrv,*ptrh);

    // If the state for the next event is not recorded, then
    // record it.
    if ( random_state.get_state_count() == itrack ) {
      int reg_flag = random_state.record();
      assert( reg_flag == itrack );
    }
    else assert( random_state.get_state_count() > itrack );

  }

  //********************************************************************
  // Write out the histogram file.  Destructor does this.
  //********************************************************************

  return 0;

}