// AddFitKalman.cpp

#include "AddFitKalman.h"
#include <iostream>
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/trfstream.h"
#include "trfbase/Hit.h"
#include "trfbase/Surface.h"
#include "HTrack.h"

using std::ostream;
using std::cout;
using std::endl;

using trf::TrfVector;
using trf::TrfMatrix;
using trf::TrfSMatrix;
using trf::HitPtr;
using trf::ETrack;
using trf::AddFitKalman;

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

namespace {

// Creator
ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "AddFitKalman" );
  if ( data.get_object_type() != "AddFitKalman" ) return ObjPtr(0);
  AddFitKalman* pobj = new AddFitKalman;
  if ( data.has("debug") ) {
    bool flag = data.get_bool("debug");
    pobj->set_debug(flag);
  }
  return ObjPtr(pobj);
}

}

//**********************************************************************
// Nested class methods.
//**********************************************************************

// Constructor.
// Allocate the space required for hit vectors, hit errors, derivatives
// and gains for hits of dimension size.

AddFitKalman::Box::Box(int size)
: _size(size) {

  int icnt;

  for ( icnt=0; icnt<NVECTOR; ++icnt ) 
    _vectors.push_back( VectorPtr(new TrfVector(size)) );

  for ( icnt=0; icnt<NERROR; ++icnt ) 
    _errors.push_back( SMatrixPtr(new TrfSMatrix(size)) );

  for ( icnt=0; icnt<NDERIV; ++icnt ) 
    _derivs.push_back( MatrixPtr(new TrfMatrix(size,5)) );

  for ( icnt=0; icnt<NDERIV; ++icnt )
    _gains.push_back( MatrixPtr(new TrfMatrix(5,size)) );

}

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

// destructor

AddFitKalman::Box::~Box() { }

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

// fetch a vector

TrfVector& AddFitKalman::Box::get_vector(int ivec) {
  assert( ivec < NVECTOR );
  return *_vectors[ivec];
}

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

// fetch an error

TrfSMatrix& AddFitKalman::Box::get_error(int ierr) {
  assert( ierr < NERROR );
  return *_errors[ierr];
}

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

// fetch a derivative (Nx5 matrix)

TrfMatrix& AddFitKalman::Box::get_deriv(int ider) {
  assert( ider < NDERIV );
  return *_derivs[ider];
}

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

// fetch a gain (5xN matrix)

TrfMatrix& AddFitKalman::Box::get_gain(int igai) {
  assert( igai < NDERIV );
  return *_gains[igai];
}

//**********************************************************************
// AddFitKalman methods.
//**********************************************************************

// Output stream.

// Fit with Kalman filter.
void AddFitKalman::ostr( ostream& stream ) const {
  stream << begin_object;
  stream << "Fit kalman";
  stream << end_object;
}

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

// Constructor.
// Allocate space needed for hits of dimension from 1 to MAXDIM.

AddFitKalman::AddFitKalman() {

  // Create boxes for hit containers.
  for ( int dim=1; dim<MAXDIM; ++dim )
    _boxes.push_back( BoxPtr(new Box(dim)) );

  int icnt;

  for ( icnt=0; icnt<MAXVECTOR; ++icnt )
    _tvectors.push_back( VectorPtr(new TrfVector(5)) );

  for ( icnt=0; icnt<MAXERROR; ++icnt )
    _terrors.push_back( SMatrixPtr(new TrfSMatrix(5)) );

}

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

// destructor

AddFitKalman::~AddFitKalman() { }

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

// Return the creator.

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

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

// Write object data.

ObjData AddFitKalman::write_data() const {
  return ObjData("AddFitKalman");
}

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

// Use Kalman filter to add a hit to a track.
// The hit is updated with the input track.
// Note: We make direct use of the underlying vector and matrix
//       classes here.  It will probably be neccessary to modify
//       this routine if these are changed.

int AddFitKalman::
add_hit_fit(ETrack& tre, double& chsq, const HitPtr& phit) const {

  // Fetch reference to the hit
  Hit& hit = *phit;

  // Update the hit with the input track.
  hit.update(tre);

  // Fetch hit size.
  int dim = hit.size();
  assert( dim <= MAXDIM );

  // Fetch the box holding the needed hit containers.
  // The chice of boxes depends on the size of the hit.
  Box& box = *_boxes[dim-1];
  assert( box.get_size() == dim );

  // Fetch the hit containers.
  TrfVector& diff    = box.get_vector(0);
  TrfVector& hit_res = box.get_vector(1);
  TrfSMatrix& hit_err     = box.get_error(0);
  TrfSMatrix& hit_err_tot = box.get_error(1);
  TrfSMatrix& hit_res_err = box.get_error(2);
  TrfMatrix& dhit_dtrk         = box.get_deriv(0);
  TrfMatrix& new_dhit_dtrk     = box.get_deriv(1);
  TrfMatrix& trk_err_dhit_dtrk = box.get_gain(0);
  TrfMatrix& gain              = box.get_gain(1);

  // Fetch the track containers.
  TrfVector& new_vec = *_tvectors[0];
  TrfSMatrix& new_err = *_terrors[0];

  hit_err = hit.measured_error().get_matrix();

  // Fetch track prediction of hit.
  diff = hit.difference_vector().get_vector();
  dhit_dtrk = hit.dhit_dtrack().get_matrix();

  // Fetch track info.
  const TrfVector& trk_vec = tre.get_vector().get_vector();
  const TrfSMatrix& trk_err = tre.get_error().get_matrix();

  // Build gain matrix.
  hit_err_tot = hit.predicted_error().get_matrix() + hit_err;
  if ( get_debug() ) {
    cout << endl;
    cout << "Begin AddFitKalman::fit(tre, chsq, phit) const";
    cout << endl;
    cout << "diff: " << endl << diff << endl;
    cout << endl;
    cout << "hit_err: " << endl << hit_err << endl;
    cout << endl;
    cout << "hit predicted error: " << endl 
         << hit.predicted_error().get_matrix() << endl;
    cout << endl;
    cout << "local hit predcted error (should be the same): " << endl 
         << dhit_dtrk % trk_err << endl;
    cout << endl;
    cout << "dhit_dtrk: " << endl << dhit_dtrk << endl;
    cout << endl;
    cout << "hit_err_tot: " << endl << hit_err_tot << endl;
  }
  if ( invert(hit_err_tot) ) return 3;
  if ( get_debug() ) {
    cout << endl;
    cout << "hit_err_tot (after inversion): " << endl
         << hit_err_tot << endl;
  }
  trk_err_dhit_dtrk = trk_err * dhit_dtrk.transpose();
  gain = trk_err_dhit_dtrk * hit_err_tot;

  if ( get_debug() ) {
    cout << endl;
    cout << "trk_vec: " << endl << trk_vec << endl;
    cout << endl;
    cout << "trk_err: " << endl << trk_err << endl;
    cout << endl;
    cout << "trk_err_dhit_dtrk: " << endl << trk_err_dhit_dtrk << endl;
  }

  // We need to return dhit_dtrk to its original state for the
  // next call.
  dhit_dtrk.transpose();

  // Build new track vector.
  new_vec = trk_vec - gain*diff;

  // Build new error;
  new_err = trk_err - trk_err_dhit_dtrk % hit_err_tot;

  if ( get_debug() ) {
    cout << endl;
    cout << "new_err: " << endl << new_err << endl;
  }

  // Check the error.
  {
    int nbad = 0;
    for ( int i=0; i<5; ++i ) {
      if ( new_err(i,i) < 0.0 ) ++nbad;
      double eii = new_err(i,i);
      for ( int j=0; j<i; ++j ) {
        double ejj = new_err(j,j);
        double eij = new_err(j,i);
        if ( fabs(eij*eij) >= eii*ejj ) ++nbad;
      }
    }
    if ( nbad > 0 ) return 5;
  }

  // Create track vector with new values.
  tre.set_vector_and_keep_direction(TrackVector(new_vec));
  tre.set_error(TrackError(new_err));

  // Calculate residual vector.

  // Update the hit with the new track.
  hit.update(tre);

  hit_res = hit.difference_vector().get_vector();
  new_dhit_dtrk = hit.dhit_dtrack().get_matrix();

  // Calculate residual covariance and invert.
  hit_res_err = hit_err - dhit_dtrk%new_err;
  if ( invert( hit_res_err ) ) return 4;
 
  // Update chi-square.
  double dchsq = hit_res_err % hit_res;
  chsq = chsq + dchsq;

  if ( get_debug() ) {
    cout << endl;
    cout << "gain: " << endl << gain << endl;
    cout << endl;
    cout << "new_vec: " << endl <<new_vec << endl;
    cout << endl;
    cout << "new_err: " << endl << new_err << endl;
    cout << endl;
    cout << "new_dhit_dtrk: " << endl << new_dhit_dtrk << endl;
    cout << endl;
    cout << "hit_res: " << endl << hit_res << endl;
    cout << endl;
    cout << "I hit_res_err: " << endl << hit_res_err << endl;
    cout << endl;
    cout << "dchsq & chsq: " << dchsq << " " << chsq << endl;
    cout << endl;
    cout << "End AddFitKalman::fit(tre, chsq, phit) const: success";
    cout << endl;
  }

  return 0;
}

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