// GTrackPropagator.cpp

#include "GTrackPropagator.hpp"
#include <cassert>
#include <iostream>
#include <map>
#include "trfbase/PropStat.h"
#include "trfbase/PropagatorPtr.h"
#include "FitStatus.hpp"
#include "GtrTrfPropagator.hpp"

using std::ostream;
using std::endl;
using std::pair;
using std::map;
using std::less;
using std::abs;

using namespace trf;

//**********************************************************************
// Local constants, typedefs, fucntions and classes.
//**********************************************************************

namespace {

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

// Order two doubles.
void order(double& x1, double& x2) {
  if ( x1 > x2  ) {
    double tmp = x1;
    x1 = x2;
    x2 = tmp;
  }
}

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

// A bad state.
GTrackState BAD_STATE;

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

// A bad S-value.
const double SBAD = 2.0*GTrack::SMAX;

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

// State list.
typedef GTrack::StateList SList;
 
// State iterator
typedef SList::const_iterator SIter;

// Pair of iterators.
typedef pair<SIter,SIter> SPair;

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

// Class to calculate and cache s for propagation from states to
// a surface.
class SCalc {
private:  // typedefs
  // Map of s indexed by iterator.
  typedef map< GTrackState, double, GTrack::StateCompare > SMap;
private: // attributes
  const SList& _states;
  const SurfacePtr& _psrf;
  const Propagator& _prop;
  SMap _smap;
public:  // methods
  // constructor
  SCalc(const SList& _states, const SurfacePtr& _psrf,
        const Propagator& pprop);
  // Return s for an iterator.
  double s(const SIter& istate);
  // Return the state list.
  const SList& states() const { return _states; }
  // Return the surface.
  const SurfacePtr& surface() const { return _psrf; }
};

// constructor
SCalc::SCalc(const GTrack::StateList& states, const SurfacePtr& psrf,
             const Propagator& prop)
: _states(states), _psrf(psrf), _prop(prop) { }

// Return s for an iterator.
double SCalc::s(const SIter& istate) {
  assert( istate != _states.end() );
  if ( istate == _states.end() ) return SBAD;
  SMap::const_iterator ipair = _smap.find(*istate);
  if ( ipair == _smap.end() ) {
    if ( istate->has_valid_fit() ) {
      VTrack trv = istate->track();
      PropStat pstat = _prop.vec_prop(trv,*_psrf);
      double s = SBAD;
      if ( pstat.success() ) s = pstat.path_distance();
      _smap[*istate] = s;
    } else {
      _smap[*istate] = SBAD;
    }
  }
  return _smap[*istate];
}

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

// Find the the states neighboring a specified surface.
// (end,end)     ==> poppagation error or out of range
// (--end,end)   ==> after end of list
// (end,begin)   ==> before beginning of list
// (ista,++ista) ==> between x and ++x (ista != end)
// (ista,ista)   ==> at ista (ista != end)
//
// We exit if we ever encounter a propagation error.
// One might want to try something more exhaustive (e.g. all
// surfaces) if this fails.
//
SPair find_neighbors(SCalc& scalc, double s1, double s2) {

  // Fetch the states and surface.
  const SList& states = scalc.states();
  const SurfacePtr& psrf = scalc.surface();

  // For convenience name the beginning and end of the list.
  const SIter& begin = states.begin();
  const SIter& end = states.end();
  SIter last = end;
  --last;

  // List may not be empty
  if ( states.size() == 0 ) return SPair(end,end);

  // Find the first state in the s-range.
  // Use the first state if there are none.
  order(s1,s2);
  SIter ista = states.lower_bound( GTrackState(s1) );
  if ( ista == end ) ista = begin;

  // Propagate from the first state to get an estimate for s.
  // Exit if propagation fails.
  double ds = scalc.s(ista);
  if ( ds == SBAD ) return SPair(end,end);
  double s = ista->s() + ds;

  // Use estimate of s to find the last state before or at the
  // surface.  If there is no such state, use the last state.
  ista = states.lower_bound( GTrackState(s) );
  if ( ista == end ) --ista;
  ds = scalc.s(ista);
  s = ista->s() + ds;
  if ( ds == SBAD ) return SPair(end,end);
  if ( ds == 0.0 ) {
    if ( s>=s1 && s<=s2 ) return SPair(ista,ista);
    else return SPair(end,end);
  }

  // Find a state after the surface.
  while ( ds > 0.0  ) {
    ++ista;
    // After last state.
    if ( ista == end ) {
      if ( s>=s1 && s<=s2 ) return SPair(last,end);
      else return SPair(end,end);
    }
    ds = scalc.s(ista);
    s = ista->s() + ds;
    if ( ds == SBAD ) return SPair(end,end);
    if ( ds == 0.0 ) {
      if ( s>=s1 && s<=s2 ) return SPair(ista,ista);
      else return SPair(end,end);
    }
  }
    
  // Find first state before the surface.
  assert( ds < 0.0 );
  while ( ds < 0.0 ) {
    if ( ista == begin ) {
      if ( s>=s1 && s<=s2 ) return SPair(end,begin);
      else return SPair(end,end);
    }
    --ista;
    ds = scalc.s(ista);
    s = ista->s() + ds;
    if ( ds == SBAD ) return SPair(end,end);
    if ( ds == 0.0 ) {
      if ( s>=s1 && s<=s2 ) return SPair(ista,ista);
      else return SPair(end,end);
    }
  }

  // Check for logic errors.
  SIter ista1 = ista;
  SIter ista2 = ista;
  ++ista2;
  assert( ista1 != ista2 );
  assert( ista2 != begin );
  assert( ista1 != end );
  assert( ista2 != end );
  assert( scalc.s(ista1) != SBAD );
  assert( scalc.s(ista2) != SBAD );
  assert( scalc.s(ista1) > 0.0 );
  assert( scalc.s(ista2) < 0.0 );

  // We have found the states bracketing our surface.
  if ( s>=s1 && s<=s2 ) return SPair(ista1,ista2);
  else return SPair(end,end);

}

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

// Propagate forward from the specified state to the specified surface
// and return a new state.
//
// from_opt is the fit status to use if the initial status is OPTIMAL.
GTrackState
prop_forward(const GTrackState& state, const SurfacePtr& psrf,
             const Propagator& prop, FitStatus from_opt =OPTIMAL) {
  ETrack tre = state.track();
  PropStat pstat = prop.err_prop(tre,*psrf);
  if ( ! pstat.forward() ) return GTrackState();
  double s = state.s() + pstat.path_distance();
  GTrackState::FitStatus fstat = state.fit_status();
  switch ( fstat ) {
  case OPTIMAL:  fstat = from_opt; break;
  case FORWARD:                    break;
  case BACKWARD: fstat = COMPLETE; break;
  case PULL:     fstat = PARTIAL;  break;
  case COMPLETE:                   break;
  case PARTIAL:                    break;
  default:       assert(false);
  }
  double chsq = state.chi_square();
  return GTrackState(s, tre, fstat, chsq);
}

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

// Propagate backward from the specified state to the specified surface
// and return a new state.
//
// from_opt is the fit status to use if the initial status is OPTIMAL.
GTrackState
prop_backward(const GTrackState& state, const SurfacePtr& psrf,
              const Propagator& prop, FitStatus from_opt =OPTIMAL) {
  ETrack tre = state.track();
  PropStat pstat = prop.err_prop(tre,*psrf);
  if ( ! pstat.backward() ) return GTrackState();
  double s = state.s() + pstat.path_distance();
  GTrackState::FitStatus fstat = state.fit_status();
  switch ( fstat ) {
  case OPTIMAL:  fstat = from_opt; break;
  case FORWARD:  fstat = COMPLETE; break;
  case BACKWARD:                   break;
  case PULL:     fstat = PARTIAL;  break;
  case COMPLETE:                   break;
  case PARTIAL:                    break;
  default:       assert(false);
  }
  double chsq = state.chi_square();
  return GTrackState(s, tre, fstat, chsq);
}

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

}  // end unnamed namespace

//**********************************************************************
// Member functions.
//**********************************************************************

// Default constructor.
// Constructs default prpagator.

GTrackPropagator::GTrackPropagator() {
  // If propagator is not defined, construct one.
  _pprop = GtrTrfPropagator().get_propagator();
}

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

// Constructor from a propagator.
// Propagator may be null.

GTrackPropagator::GTrackPropagator(const PropagatorPtr& pprop)
: _pprop(pprop) { }

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

// Destructor.

GTrackPropagator::~GTrackPropagator() { }

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

// Propagate to a state.

GTrackState GTrackPropagator::
propagate(const GTrack& gtr, const SurfacePtr& psrf,
          double s1, double s2) const {

  // If the propagator is null, then only search the states which
  // already appear in the list.
  if ( ! _pprop ) return gtr.get_state(psrf,s1,s2);

  // construct the s calculator
  SCalc scalc( gtr.get_states(), psrf, *_pprop );

  // For convenience name the beginning and end of the list.
  const SIter& begin = gtr.get_states().begin();
  const SIter& end = gtr.get_states().end();
  SIter last = end;
  --last;

  // Find the neighboring states.
  SPair spair = find_neighbors(scalc,s1,s2);
  SIter ista1 = spair.first;
  SIter ista2 = spair.second;
  SIter ista1_plus = ista1;
  if ( ista1 != end ) ++ista1_plus;

  // Propagation failure.
  if ( ista1==end && ista2==end ) {
    return GTrackState();
  }

  // Before the beginning of the list.
  // If fit was optimal, it remains so.
  else if ( ista1==end && ista2==begin ) {
    return prop_backward(*ista2,psrf,*_pprop);
  }

  // After the end of the list.
  // If fit was optimal, it remains so.
  else if ( ista1==last && ista2==end ) {
    return prop_forward(*ista1,psrf,*_pprop);
  }

  // At a surface.
  else if ( ista1!=end && ista1==ista2 ) {
    return *ista1;
  }

  // Between two surfaces.
  // Here is where we could use smoothing.
  // For now extrapolate from closest surface.
  else if ( ista1!=end && ista1_plus==ista2 ) {
    if ( abs(scalc.s(ista1)) <= abs(scalc.s(ista2)) ) {
      return prop_forward(*ista1,psrf,*_pprop,COMPLETE);
    } else {
      return prop_backward(*ista2,psrf,*_pprop,COMPLETE);
    }
  }

  // Error - unexpected return value.
  else {
    return GTrackState();
  }

}

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

// Output stream.

ostream& operator<<(ostream& stream, const GTrackPropagator& rhs) {
  stream << "GTrack propagator using TRF++ propagator:" << endl;
  stream << *rhs.get_propagator();
  return stream;
}

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

// Equality.

bool
operator==(const GTrackPropagator& lhs, const GTrackPropagator& rhs) {
  return lhs.get_propagator() == rhs.get_propagator();
}

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