// ClusFindZPlane2.cpp

#include "ClusFindZPlane2.h"
#include "HitZPlane2.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() == "ClusFindZPlane2" );
  double z = data.get_double("z");
  double diff = data.get_double("max_chsq_diff");
  return ObjPtr( new ClusFindZPlane2( z,diff) );
}

}

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

// Calculate the chi-square difference between and cluster and a track.
// Only x components are taken into account for calculated returned chisq_diff
// Real chisquare are returned in double *chisq variable
// It is assumed the cluster generates exactly one prediction and that
// this prediction has two dimension.

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

  const ClusZPlane2& clu = (const ClusZPlane2&)(*pclu);
  const TrackVector& vec = tre.get_vector();
  const TrackError& err = tre.get_error();

  double diff0 = clu.get_x() - vec(SurfZPlane::IX ); 
  double diff1 = clu.get_y() - vec(SurfZPlane::IY ); 

  if ( diff0 == 0. && diff1 == 0. ) return 0.;

  double e11 = err(SurfZPlane::IX,SurfZPlane::IX);
  double e22 = err(SurfZPlane::IY,SurfZPlane::IY);
  double e12 = err(SurfZPlane::IY,SurfZPlane::IX);
  double a = e11 + clu.get_dx2();
  double b = e12 + clu.get_dxdy();
  double c = e22 + clu.get_dy2();
  double det = a*c-b*b;
 
  if ( det  == 0. ) return 1.e15;

  double chisq = (c*diff0*diff0 + a*diff1*diff1 - 2*b*diff0*diff1)/det; 

  if ( chisq < 0 ) return 1.e15;
  
  return chisq;
  
}

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

// ouput stream

void ClusFindZPlane2::ostr(ostream& stream) const {
  stream << begin_object ;
  stream << "2D Cluster finder for " << _szp << "." << new_line;
  stream << "Maximum chi-square difference is " 
         << _max_chsq_diff << ".";
  stream << end_object;
}

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

// constructor

ClusFindZPlane2::
ClusFindZPlane2(double zpos, double max_chsq_diff)
: _szp(zpos), _max_chsq_diff(max_chsq_diff) { }

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

// destructor

ClusFindZPlane2::~ClusFindZPlane2() { }

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

// add a cluster

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

  // Extract the cluster position depending on type
  if ( pclu->get_type() == ClusZPlane2::get_static_type() ) {
    const ClusZPlane2& clu = (const ClusZPlane2&) *pclu;
    //assert( clu.get_surface() == _szp );
  }  else {
    assert(false);
    return 1;
  }

  const ClusterList& clusters = _box.get_clusters();
  // Check there is no other cluster at this position.
  ClusterList::const_iterator iclu ;
  for ( iclu= clusters.begin() ;  iclu!= clusters.end(); ++iclu )
    assert( (*iclu) != pclu );

  // Store the cluster.
  _box.add_cluster(&*pclu);

  return 0;

}

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

// drop all clusters

void ClusFindZPlane2::drop_clusters() {

  _box.drop_clusters();

}

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

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

  return _box.get_clusters();

}

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

const ClusterContainer& ClusFindZPlane2::get_cluster_container(
			 const ETrack& tre, const CutRecord* pcut_record,
                         std::map< ClusterPtr, CutRecord*>* precs) const {

  if( _max_chsq_diff >= 0. ) {
    fill_container(tre);
    return _outbox;
  }
 
  return _box;
}

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

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

ClusterList ClusFindZPlane2::get_clusters(const ETrack& tre) const {
  if ( _max_chsq_diff < 0. ) return _box.get_clusters();
  fill_container(tre);
  return _outbox.get_clusters();
}

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

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

void ClusFindZPlane2::fill_container(const ETrack& tre) const {

   // Create the list of clusters.
  _outbox.drop_clusters();
  
  // Check the track surface.
  if ( *tre.get_surface() != _szp ) {
    assert(false);
    return;
  }
  
  const ClusterList& _clusters = _box.get_clusters();

  // If there are no clusters, return the empty list.
  if ( ! _clusters.size() ) return;
  
  ClusterList::const_iterator iclu;
  for(iclu = _clusters.begin() ; iclu != _clusters.end(); ++iclu ) {
    const ClusterPtr& pclu = (*iclu);
    if ( chsq_diffzp2(pclu,tre) <= _max_chsq_diff ) _outbox.add_cluster( &*pclu );
  }

}

//**********************************************************************
// Return the creator.

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

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

// Write the object data.

ObjData ClusFindZPlane2::write_data() const {
  ObjData data( get_type_name() );
  data.add_double("z", _szp.get_z() );
  data.add_double( "max_chsq_diff", get_max_chsq_diff() );
  return data;
}

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