// ClusFindZPlane1.cpp

#include "ClusFindZPlane1.h"
#include "HitZPlane1.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() == "ClusFindZPlane1" );
  double z = data.get_double("z");
  double wx = data.get_double("xweight");
  double wy = data.get_double("yweight");
  double diff = data.get_double("max_chsq_diff");
  return ObjPtr( new ClusFindZPlane1( z, wx,wy, 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_diffzp1(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 ClusFindZPlane1::ostr(ostream& stream) const {
  stream << begin_object ;
  stream << "1D Cluster finder for " << _szp << "." << new_line;
  stream << "wx weight is " << _wx << "." << new_line;
  stream << "wy weight is " << _wy << "." << new_line;
  stream << "Maximum chi-square difference is " 
         << _max_chsq_diff << ".";
  stream << end_object;
}

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

// constructor

ClusFindZPlane1::
ClusFindZPlane1(double zpos, double wx, double wy, double max_chsq_diff)
: _szp(zpos), _wx(wx), _wy(wy), _max_chsq_diff(max_chsq_diff) { }

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

// destructor

ClusFindZPlane1::~ClusFindZPlane1() { }

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

// add a cluster

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

  // declare the key value
  double keyval;

  // Extract the cluster position depending on type
  if ( pclu->get_type() == ClusZPlane1::get_static_type() ) {
    const ClusZPlane1& clu = (const ClusZPlane1&) *pclu;
    assert( clu.get_surface() == _szp );
    assert( clu.get_wx() == _wx );
    assert( clu.get_wy() == _wy );
    keyval = clu.get_axy();
  }  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 ClusFindZPlane1::drop_clusters() {

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

}

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

// Return all clusters.
ClusterList ClusFindZPlane1::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 axy
// The nearby subset is returned in the same order.

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

  // Create the list of clusters.
  ClusterList clusters;

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

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

  // Fetch the predicted position.
  double value = _wx*tre.get_vector(SurfZPlane::IX) + 
                 _wy*tre.get_vector(SurfZPlane::IY);

  // 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.

  ClusterMap::const_iterator iclu = iclu1;
  for(iclu = iclu1 ; iclu != _clusters.end(); ++iclu ) {
    ClusterPtr pclu = (*iclu).second;
    if ( chsq_diffzp1(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_diffzp1(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 ClusFindZPlane1::get_creator() {
  return create;
}

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

// Write the object data.

ObjData ClusFindZPlane1::write_data() const {
  ObjData data( get_type_name() );
  data.add_double("z", _szp.get_z() );
  data.add_double( "xweight", _wx );
  data.add_double( "yweight", _wy );
  data.add_double( "max_chsq_diff", get_max_chsq_diff() );
  return data;
}

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