// AddFitCyl_PhiZ_PhiZ.cpp

#include <iostream>
#include "AddFitCyl_PhiZ_PhiZ.h"
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/trfstream.h"
#include "trfutil/TRFMath.h"
#include "trfbase/Propagator.h"
#include "trfbase/PropStat.h"
#include "trffit/HTrack.h"
#include "HitCylPhi.h"
#include "HitCylPhiZ.h"

using std::vector;
using std::cerr;
using std::endl;

using namespace trf;

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

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

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "AddFitCyl_PhiZ_PhiZ" );
  ObjDoublePtr bfield;
  ObjData::Type dtype = data.get_type("bfield");
  if(dtype == ObjData::DOUBLE)
    bfield = new ObjDouble(data.get_double("bfield"));
  else if(dtype == ObjData::SHARE_PTR)
    bfield.assign_with_dynamic_cast(data.get_share_pointer("bfield"));
  else {
    cerr << "AddFitCyl_PhiZ_PhiZ: can't get bfield parameter." << endl;
    abort();
  }
  const AddFitter* pfit;
  data.get_bare_pointer("fitter",pfit);
  const Propagator* pprop;
  data.get_bare_pointer("propagator",pprop);
  SurfacePtr psrf;
  data.get_share_pointer("surface",psrf);
  TrackError err;
#ifdef ObjData_supports_lists
  vector<double> vals;
  data.get_double_list("error_matrix",vals);
  assert( vals.size() == 15 );
  int ival = 0;
  for ( int j=0; j<5; ++j ) {
    for ( int i=0; i<=j; ++i ) {
      err(i,j) = vals[ival++];
    }
  }
#endif
  double z0max = data.get_double("z0max");
  double z0min = -z0max;
  if ( data.has("z0min") ) {
    z0min = data.get_double("z0min");
  }
  string name = data.get_string("flag");
  AddFitCyl_PhiZ_PhiZ::FlagState flag(AddFitCyl_PhiZ_PhiZ::DR_ALPHA);
  if ( name == "DR_ALPHA" ) flag = AddFitCyl_PhiZ_PhiZ::DR_ALPHA;
  else if ( name == "DR_ONLY" ) flag = AddFitCyl_PhiZ_PhiZ::DR_ONLY;
  else if ( name == "DR_ALPHA_CURV" ) flag = AddFitCyl_PhiZ_PhiZ::DR_ALPHA_CURV;
  else assert(false);

  int fit_order=AddFitCyl_PhiZ_PhiZ::UNCHANGED;

  if( data.has("fit_order") ) {
    string fitorder= data.get_string("fit_order");
    if( fitorder == "INWARD" ) fit_order=AddFitCyl_PhiZ_PhiZ::INWARD;
    else if( fitorder == "OUTWARD" ) fit_order=AddFitCyl_PhiZ_PhiZ::OUTWARD;
    else if( fitorder != "UNCHANGED" ) assert(false);
  }

  return ObjPtr( new 
    AddFitCyl_PhiZ_PhiZ(bfield,*pfit,*pprop,psrf,err,z0min,z0max,flag,fit_order) );
}

}

//**********************************************************************
// local helper functions
//**********************************************************************

// Use a HitCylPhiZ object and a phi value to calculate r and z.
// Return 0 for no errors.

static int get_r_z(const Hit& hit, double phi, double& r, double& z) {

  assert( hit.get_type() == HitCylPhiZ::get_static_type() );
  // Fetch r from the surface.
  r = hit.get_surface().get_parameter(0);
  
  // Fetch the stereo parameter and measurement.
  double stereo = hit.dhit_dtrack()(0,1);
  if ( stereo == 0.0 ) return 1;
  double phiz = hit.measured_vector()(0);

  // Calculate z.
  // phiz = phi +stereo*z
  double z0 = (phiz-phi)/stereo;

  // Since phiz is cyclic, there are multiple solutions for z.
  // Choose the one closest to zero.
  z = fmod2( z0, fabs(TWOPI/stereo) );

  return 0;

}

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

// ouput stream

void AddFitCyl_PhiZ_PhiZ::ostr(ostream& stream) const {
  stream << begin_object;
  stream << "Add fitter for two cylinder phiz measurements." << new_line;
  stream << "Internal fitter is " << _fit << new_line;
  stream << "Internal propagator is " << _prop << new_line;
  stream << "Starting surface is " << *_psrf << new_line;
  stream << "Starting error matrix is" << new_line << _err << new_line;
  stream << "Z0 range is (" << _z0min << ", " << _z0max << ")";
  stream << end_object;
}

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

// constructor

AddFitCyl_PhiZ_PhiZ::
AddFitCyl_PhiZ_PhiZ(ObjDoublePtr bfield, const AddFitter& fit,
                    const Propagator& prop, const SurfacePtr& psrf,
                    const TrackError& err, double z0min, double z0max,
                    FlagState flag)
: _bfield(bfield),
  _fit(fit),
  _prop(prop),
  _psrf(psrf),
  _err(err), 
  _flag(flag),
  _fit_order(UNCHANGED),
  _z0min(z0min), _z0max(z0max) { }

AddFitCyl_PhiZ_PhiZ::
AddFitCyl_PhiZ_PhiZ(ObjDoublePtr bfield, const AddFitter& fit,
                    const Propagator& prop, const SurfacePtr& psrf,
                    const TrackError& err, double z0min, double z0max,
                    FlagState flag, int fit_order)
: _bfield(bfield),
  _fit(fit),
  _prop(prop),
  _psrf(psrf),
  _err(err), 
  _flag(flag),
  _fit_order(fit_order),
  _z0min(z0min), _z0max(z0max) { }

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

// destructor

AddFitCyl_PhiZ_PhiZ::~AddFitCyl_PhiZ_PhiZ() { }

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

// Return the creator.

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

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

// Write the object data.

ObjData AddFitCyl_PhiZ_PhiZ::write_data() const {
  ObjData data( get_type_name() );
  data.add_double( "bfield", get_bfield() );
  data.add_bare_pointer( "fitter", &get_fitter() );
  data.add_bare_pointer( "propagator", &get_propagator() );
  data.add_share_pointer( "surface", get_surface() );
#ifdef ObjData_supports_lists
  const TrackError& err = get_error();
  vector<double> vals;
  for ( int j=0; j<5; ++j ) {
    for ( int i=0; i<=j; ++i ) {
      vals.push_back( err(i,j) );
    }
  }
  data.add_double_list( "error_matrix", vals );
#endif
  data.add_double("z0min", _z0min);
  data.add_double("z0max", _z0max);
  string name;
  switch ( get_flag() ) {
  case DR_ONLY:
    name="DR_ONLY"; break;
  case DR_ALPHA:
    name="DR_ALPHA"; break;
  case DR_ALPHA_CURV:
    name="DR_ALPHA_CURV"; break;
  default:
    assert(false);
  }
  data.add_string( "flag", name );

  if( _fit_order != UNCHANGED ) {
    string fitorder="OUTWARD";
    if ( _fit_order == INWARD ) fitorder="INWARD";
    data.add_string("fit_order",fitorder);
  }
  return data;
}

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

// add the hit

int AddFitCyl_PhiZ_PhiZ::
add_hit(HTrack& trh, const HitPtr& phit) const {

  // check type
  if ( phit->get_type() != HitCylPhiZ::get_static_type() ) return 1;

  // Loop over hits: there should be one HitCylPhiZ and any number of
  // HitCylPhi's.
  HitList::const_iterator ihit;
  HitList hits = trh.get_hits();
  int nphiz = 0;
  int nphi = 0;
  const HitCylPhiZ* phit1 = 0;
  for ( ihit=hits.begin(); ihit!=hits.end(); ++ihit ) {
    Hit& hit = **ihit;
    if ( hit.get_type() == HitCylPhi::get_static_type() ) ++nphi;
    else
      if ( hit.get_type() == HitCylPhiZ::get_static_type() ) {
        ++nphiz;
        phit1 = (const HitCylPhiZ*) &hit;
      }
      else return 2;
  }  // end loop over hits
  if ( nphiz != 1 ) return 3;

  // Add the new hit to the internal list of hits.
  hits.push_back(phit);

  // Fetch the track vector and error.
  VTrack trv2 = trh.get_track();
  VTrack trv = trv2;

  // Fetch r and z for the second hit.
  // Assume lambda = 0.0.
  double bfac = -BFAC * (*_bfield)();
  double phi2 = trv.get_vector(IPHI);
  double alf2 = trv.get_vector(IALF);
  double qpt2 = trv.get_vector(IQPT);
  double crv2 = bfac*qpt2;
  double r2, z2;
  int stat2 = get_r_z( *phit, phi2, r2, z2 );
  if ( stat2 ) return 3+stat2;

  // Propagate track to first phiz hit.
  const Surface& srf1 = phit1->get_surface();
  PropStat pstat = _prop.vec_prop(trv,srf1);
  if ( ! pstat.success() ) return 5;

  // Fetch r and z for the first hit.
  // Assume lambda = 0.0.
  double phi1 = trv.get_vector(IPHI);
  double alf1 = trv.get_vector(IALF);
  double qpt1 = trv.get_vector(IQPT);
  double crv1 = bfac*qpt1;
  double r1, z1;
  int stat1 = get_r_z( *phit1, phi1, r1, z1 );
  if ( stat1 ) return 5+stat1;
   
  // check radii
  if ( r1 == r2 ) return 7;

  // Begin estimate for tan(lambda).
  double dr = r2 - r1;
  double alf = 0.5*(alf1+alf2);
  double crv = 0.5*(crv1+crv2);

  // Calculate dsT (transverse path length) based on _flag.
  double dst = dr;
  double st2 = r2;
  double calf = cos(alf);
  switch ( _flag ) {
  // Use sT = r2 - r1 (assume alpha = C = 0)
  case DR_ONLY:
    break;
  // Use sT = (r2-r1)/cos(alpha) (assume C = 0)
  case DR_ALPHA:
    dst = dr/calf;
    st2 = st2/calf;
    break;
  // Use sT = x*asin(xC/2)/(xC/2) where x = (r2-r1)/cos(alpha)
  case DR_ALPHA_CURV:
    dst = dr/cos(alf);
    dst = dst*asinrat(0.5*dst*crv);
    st2 = st2/calf;
    st2 = st2*asinrat(0.5*st2*crv);
    break;
  // unknown flag
  default:
    assert(false);
    return 8;
  }

  // Calculate tan(lambda) = dz/dsT
  double tlam = (z2-z1)/dst;

  // Calculate z0.
  // If track is out of range, exit with error.
  double z0 = z2 - tlam*st2;
  if ( z0 < _z0min ) {
    return 51;
  }
  if ( z0 > _z0max ) {
    return 52;
  }

  // Update the track vector with this value and propagate back
  // to the starting surface.
  TrackVector vec2 = trv2.get_vector();
  vec2(IZ) = z2;
  vec2(ITLM) = tlam;
  trv.set_vector(vec2);
  pstat = _prop.vec_prop(trv,*_psrf);
  if ( ! pstat.success() ) return 9;

  // construct and fit a new track
  ETrack tre( _psrf, trv.get_vector(), _err );
  double chsq = 0.0;

  // Loop over hits and add each in turn to the track.
  int count = 0;

  HitList nhits ;

  if( _fit_order != UNCHANGED ) {
    int size = hits.size();
    std::vector<double> rarray(size);
    std::vector<HitPtr> harray(size);
    int ii=0;
    for(ihit=hits.begin() ; ihit!=hits.end(); ihit++ )
      {
	const HitPtr& tmphit = *ihit;
	double rrr = tmphit->get_surface().get_parameter(SurfCylinder::RADIUS);
	harray[ii] = tmphit;
	rarray[ii] = rrr;
	ii++;
      }
    
    for(int i=0;i<size;++i)
      for(int j=i+1;j<size;++j)
	{
	  if( rarray[j] > rarray[i] ) {
	    HitPtr tmp = harray[j];
	    double rtmp = rarray[j];
	    harray[j] = harray[i];
	    rarray[j] = rarray[i];
	    harray[i] = tmp;
	    rarray[i] = rtmp;
	  }
	}
    
    std::vector<HitPtr>::const_iterator iihit;
    std::vector<HitPtr>::reverse_iterator riihit;
    if( _fit_order == INWARD ) 
      for(iihit=harray.begin() ; iihit!=harray.end(); iihit++ ) 
	nhits.push_back(*iihit);
    else if( _fit_order == OUTWARD ) 
      for(riihit=harray.rbegin() ; riihit!=harray.rend(); riihit++ ) 
	nhits.push_back(*riihit);
    else assert(false);
  }
  else {
    nhits = hits;
  }

  for ( ihit=nhits.begin(); ihit!=nhits.end(); ++ihit ) {
    ++count;
    const Surface& srf = (*ihit)->get_surface();
    PropStat pstat =  _prop.err_prop(tre,srf);
    if ( ! pstat.success() ) return 10+count;
    int stat = _fit.add_hit_fit(tre,chsq,*ihit);
    if ( stat ) return 100*count + stat;
  }
  
  // Fit was successful: update fit and add new hit.
  trh.set_fit(tre,chsq);
  trh.add_hit(phit);

  return 0;

}

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