// ClusFindXYPlane1.cpp

#include "ClusFindXYPlane1.h"
#include "HitXYPlane1.h"
#include "trfutil/TRFMath.h"
#include "trfbase/ETrack.h"
#include "trfutil/trfstream.h"
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"

using namespace trf;

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

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "ClusFindXYPlane1" );
  double dist = data.get_double("dist");
  double phi = data.get_double("phi");
  double wv = data.get_double("vweight");
  double wz = data.get_double("zweight");
  double diff = data.get_double("max_chsq_diff");
  return ObjPtr( new ClusFindXYPlane1( dist,phi, wv,wz, diff) );
}

}

//**********************************************************************
// Helper functions.
//**********************************************************************

// Calculate the chi-square difference between and cluster and a track.
// It is assumed the cluster generates exactly one prediction and that
// this prediction has one dimension.

double chsq_diffxyp1(const ClusterPtr& pclu, const ETrack& tre) {

  // Fetch the hit predictions.
  const HitList& hits = pclu->predict(tre);

  // Check there is one hit.
  assert( hits.size() == 1 );

  // Fetch the hit.
  const Hit& hit = *hits.front();

  // Check the hit dimension.
  assert( hit.size() == 1 );

  // Fetch the difference between prediction and measurement.
  double diff = hit.difference_vector()(0);

  // Fetch the measurement error and the prediction error.
  double emeas = hit.measured_error()(0,0);
  double epred = hit.predicted_error()(0,0);

  // Calculate and return chi-square difference.
  return diff*diff/(emeas+epred);

}

//**********************************************************************
// Methods.
//**********************************************************************

// ouput stream

void ClusFindXYPlane1::ostr(ostream& stream) const {

  stream << begin_object ;
  stream << "1D Cluster finder for " << _sxyp << "." << new_line;
  stream << "wv weight is " << _wv << "." << new_line;
  stream << "wz weight is " << _wz << "." << new_line;
  stream << "Maximum chi-square difference is " 
         << _max_chsq_diff << ".";
  stream << end_object;

}

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

// constructor

ClusFindXYPlane1::
ClusFindXYPlane1(double dist, double phi, 
                 double wv, double wz, double max_chsq_diff)
: _sxyp(dist,phi), _wv(wv), _wz(wz), _max_chsq_diff(max_chsq_diff) { }

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

// destructor

ClusFindXYPlane1::~ClusFindXYPlane1() { }

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

// add a cluster

int ClusFindXYPlane1::add_cluster(const ClusterPtr& pclu) {

  // declare the key value
  double keyval;

  // Extract the cluster position depending on type
  if ( pclu->get_type() == ClusXYPlane1::get_static_type() ) {
    const ClusXYPlane1& clu = (const ClusXYPlane1&) *pclu;
    assert( clu.get_surface() == _sxyp );
    assert( clu.get_wv() == _wv );
    assert( clu.get_wz() == _wz );
    keyval = clu.get_avz();
  }  else {
    assert(false);
    return 1;
  }

  // Check there is no other cluster at this position.
  assert( _clusters.find(keyval) == _clusters.end() );

  // Store the cluster.
  _clusters[keyval] = pclu;

  return 0;

}

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

// drop all clusters

void ClusFindXYPlane1::drop_clusters() {

  _clusters.erase( _clusters.begin(), _clusters.end() );

}

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

// Return all clusters.
ClusterList ClusFindXYPlane1::get_clusters() const {

  // Create the list of clusters.
  ClusterList clusters;

  // Loop over clusters and add to output list.
  ClusterMap::const_iterator iclu;
  for ( iclu=_clusters.begin();  iclu!=_clusters.end(); ++iclu )
    clusters.push_back( (*iclu).second );

  return clusters;

}

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

// Return the clusters near a track.
//
// Clusters have been ordered by avz
// The nearby subset is returned in the same order.

ClusterList ClusFindXYPlane1::get_clusters(const ETrack& tre) const {

  // Create the list of clusters.
  ClusterList clusters;

  // Check the track surface.
  if ( *tre.get_surface() != _sxyp ) {
    assert(false);
    return clusters;
  }

  // If there are no clusters, return the empty list.
  if ( ! _clusters.size() ) return clusters;

  // Fetch the predicted position.
  double value = _wv*tre.get_vector(SurfXYPlane::IV) + 
                 _wz*tre.get_vector(SurfXYPlane::IZ);

  // Fetch first cluster after the prediction.

  ClusterMap::const_iterator iclu1 = _clusters.lower_bound(value);

  // Loop over clusters from this and above.
  // If cluster is close, add it to the list; otherwise exit loop.
  // The last cluster checked is iclu2.
  ClusterMap::const_iterator iclu = iclu1;
  for(iclu = iclu1 ; iclu != _clusters.end(); ++iclu ) {
    ClusterPtr pclu = (*iclu).second;
    if ( chsq_diffxyp1(pclu,tre) > _max_chsq_diff ) break;
    clusters.push_back( pclu );
  }

  // We have checked all clusters in the range (iclu1, +).
  // Next go backwards from iclu1.
  iclu = iclu1;
  if( iclu == _clusters.begin() ) return clusters;
  iclu--;
  while ( true ) {
    ClusterPtr pclu = (*iclu).second;
    if ( chsq_diffxyp1(pclu,tre) > _max_chsq_diff ) break;
    clusters.push_front( pclu );
    // decrement iterator; check if the begining is reached
    if ( iclu == _clusters.begin() ) break;
    iclu--;
  }

  // Return the nearby clusters.
  return clusters;

}

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

// Return the creator.

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

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

// Write the object data.

ObjData ClusFindXYPlane1::write_data() const {
  ObjData data( get_type_name() );
  data.add_double("dist", _sxyp.get_parameter(SurfXYPlane::DISTNORM) );
  data.add_double("phi", _sxyp.get_parameter(SurfXYPlane::NORMPHI) );
  data.add_double( "vweight", _wv );
  data.add_double( "zweight", _wz );
  data.add_double( "max_chsq_diff", get_max_chsq_diff() );
  return data;
}

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