// AddFitStarter.cpp

#include "AddFitStarter.h"
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/trfstream.h"
#include "trfbase/Hit.h"
#include "trffit/HTrack.h"

using std::vector;

using trf::AddFitStarter;
using trf::AddFitterPtr;

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

namespace {

// Creator

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "AddFitStarter" );
  if ( data.get_object_type() != "AddFitStarter" ) return ObjPtr(0);
  int max_hits = data.get_int("max_hits");
  AddFitterPtr pfit0;
  data.get_share_pointer("default_fitter",pfit0);
  AddFitStarter* pfit = new AddFitStarter(max_hits);
  pfit->register_default_fitter(pfit0);
#ifdef ObjData_supports_lists
  // Read the list of registered fitters.
  vector<AddFitterPtr> rfits;
  data.get_share_ptr_list("fitters",rfits);
  // Read the list of names of types.
  vector<string> tnames;
  data.get_string_list("hit_types",tnames);
  // Loop over registered fitters.
  vector<string>::const_iterator itname = tnames.begin();
  for ( vector<AddFitterPtr>::const_iterator irfit=rfits.begin();
        irfit!=rfits.end(); ++irfit ) {
    AddFitStarter::TypeList types;
    assert( itname != tnames.end() );
    while ( *itname != ":" ) {
      assert( itname != tnames.end() );
      AddFitStarter::Type type = ObjTable::get_creator(*itname++);
      assert( type != 0 );
      types.push_back(type);
    }
    pfit->register_fitter(*irfit,types);
    ++itname;
  }
#endif
  return ObjPtr(pfit);
}

// Null fitter pointer for error return.
AddFitterPtr null_fitter(0);

}

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

// constructor

AddFitStarter::AddFitStarter(int max_hits) : _max_hits(max_hits) { }

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

// destructor

AddFitStarter::~AddFitStarter() { }

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

// Return the creator.

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

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

// Write object data.
// We need to add the list of registered fitters.

ObjData AddFitStarter::write_data() const {
  ObjData data("AddFitStarter");
  data.add_int("max_hits", _max_hits);
  data.add_share_pointer("default_fitter",_pfitter0);
#ifdef ObjData_supports_lists
  // Build the lists of fitters and hit type names.
  vector<AddFitterPtr> rfits;
  vector<string> tnames;
  // Loop over fitter map.
  for ( AddFitterMap::const_iterator ifit=_fitters.begin();
        ifit!=_fitters.end(); ++ifit ) {
    rfits.push_back(ifit->second);
    const TypeList& types = ifit->first;
    for ( TypeList::const_iterator ityp=types.begin();
          ityp !=types.end(); ++ityp )
      tnames.push_back( ObjTable::get_type_name(*ityp) );
    tnames.push_back(":");
  }
  data.add_share_ptr_list("fitters",rfits);
  data.add_string_list("hit_types",tnames);
#endif
  return data;
}

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

// set the default fitter
// return true if an old fitter has been replaced

bool AddFitStarter::
register_default_fitter(const AddFitterPtr pfitter) {
  
  // set flag indicating if default was already registered
  bool overwrite = _pfitter0.pointer() != 0;

  // register
  _pfitter0 = pfitter;

  // return the oerwrite status
  return overwrite;

}

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

// register a fitter for a list of types

bool AddFitStarter::
register_fitter(const AddFitterPtr pfitter, TypeList types) {

  // Check number of types.
  assert( types.size() <= _max_hits );
  if ( types.size() > _max_hits ) return true;

  // set flag indicating if list was already registered
  bool overwrite = _fitters.find(types) != _fitters.end();

  // register
  _fitters[types] = pfitter;

  // return the overwrite status
  return overwrite;

}

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

// register a fitter for one type

bool AddFitStarter::
register_fitter(const AddFitterPtr fitter, Type type1) {
 
  // construct the list of types
  TypeList types;
  types.push_back(type1);
 
  // register
  return register_fitter(fitter,types);

}

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

// register a fitter for two types

bool AddFitStarter::
register_fitter(const AddFitterPtr fitter,
                Type type1, Type type2) {
 
  // construct the list of types
  TypeList types;
  types.push_back(type1);
  types.push_back(type2);
 
  // register
  return register_fitter(fitter,types);

}

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

// register a fitter for three types

bool AddFitStarter::
register_fitter(const AddFitterPtr fitter,
                Type type1, Type type2, Type type3) {
 
  // construct the list of types
  TypeList types;
  types.push_back(type1);
  types.push_back(type2);
  types.push_back(type3);
 
  // register
  return register_fitter(fitter,types);

}

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

// register a fitter for four types

bool AddFitStarter::
register_fitter(const AddFitterPtr fitter,
                Type type1, Type type2, Type type3, Type type4) {
 
  // construct the list of types
  TypeList types;
  types.push_back(type1);
  types.push_back(type2);
  types.push_back(type3);
  types.push_back(type4);
 
  // register
  return register_fitter(fitter,types);

}

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

// register a fitter for five types

bool AddFitStarter::
register_fitter(const AddFitterPtr fitter, Type type1, Type type2,
                Type type3, Type type4, Type type5) {
 
  // construct the list of types
  TypeList types;
  types.push_back(type1);
  types.push_back(type2);
  types.push_back(type3);
  types.push_back(type4);
  types.push_back(type5);
 
  // register
  return register_fitter(fitter,types);

}

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

// Return the appropriate fitter for a list of hit types.

const AddFitterPtr&
AddFitStarter::get_fitter(const TypeList& types) const {
  if ( types.size() == 0  || types.size() > _max_hits )
    return _pfitter0;
  // Find fitter in registry.
  AddFitterMap::const_iterator ifit = _fitters.find(types);
  if ( ifit == _fitters.end() ) return null_fitter;
  return ifit->second;
}
  
//**********************************************************************

// add a hit
// call the registered fitter

int AddFitStarter::
add_hit_with_record(HTrack& trh, const HitPtr& phit,
                    CutRecord* prec ) const {

  // fetch # hits (plus one for the new hit)
  int nhits = trh.get_hits().size() + 1;

  // if # hits is above threshold, use default
  if ( nhits > _max_hits ) {
    if ( ! _pfitter0.pointer() ) return DEFAULT_NOT_REGISTERED;
    return _pfitter0->add_hit_with_record(trh,phit,prec);
  }

  // Build the list of types.
  TypeList types;
  HitList::const_iterator ihit;
  const HitList& hits = trh.get_hits();
  for ( ihit=hits.begin(); ihit!=hits.end(); ++ihit )
    types.push_back( (*ihit)->get_type() );
  types.push_back( phit->get_type() );

  // Find fitter in registry.
  AddFitterMap::const_iterator ifit = _fitters.find(types);
  if ( ifit == _fitters.end() ) return TYPES_NOT_FOUND;
  AddFitterPtr pfitter = (*ifit).second;

  // If zero is registered, exit.
  if ( ! pfitter.pointer() ) return TYPES_NOT_REGISTERED;

  // fit
  return pfitter->add_hit_with_record(trh,phit,prec);

}

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

// output stream

void AddFitStarter::ostr(ostream& stream) const {
  stream << begin_object;
  stream << "AddFitStarter with hit threshold " << _max_hits << "." 
         << new_line;
  if ( _pfitter0.pointer() )
    stream << "Default fitter: " << *_pfitter0;
  else stream << "No default fitter.";
  stream << new_line;
  AddFitterMap::const_iterator ifit;
  stream << "Add fitters indexed by hit types:";
  for ( ifit=_fitters.begin(); ifit!=_fitters.end(); ++ifit ) {
    const TypeList& types = (*ifit).first;
    const AddFitter* pfitter = (*ifit).second.pointer();
    stream << new_line;
    stream << "Types:";
    TypeList::const_iterator ityp;
    for ( ityp=types.begin(); ityp!=types.end(); ++ityp ) {
      stream << " " << ObjTable::get_type_name(*ityp);
    }
    stream << new_line;
    if ( pfitter ) stream << *pfitter;
    else stream << "AddFitter undefined.";
  }  // end loop over fitters
  stream << end_object;
}

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