// STL
#include <iostream>
#include <iomanip>
#include <cmath>
#include <algorithm>

// LCIO
#include "EVENT/LCIO.h"
#include "EVENT/LCCollection.h"
#include "EVENT/TrackerHit.h"
#include <IMPL/LCCollectionVec.h>
#include "IMPL/TrackImpl.h"
#include <IMPL/LCFlagImpl.h>

// Marlin
#include "marlin/Global.h"
#include "TrackFitterKalmanProcessor.h"
#include "TrackFitterBase.h"

// Aida
#ifdef MARLIN_USE_AIDA

#include <AIDA/AIDA.h>
#include <marlin/AIDAProcessor.h>

//marlin
#include <marlin/ConditionsProcessor.h>

#endif

// ROOT
#if 1
#include "TROOT.h"
#include "TNtupleD.h"
#endif

using namespace marlin;
using namespace lcio;

namespace marlintpc
{

TrackFitterKalmanProcessor aTrackFitterKalmanProcessor;

TrackFitterKalmanProcessor::TrackFitterKalmanProcessor() 
  : Processor("TrackFitterKalmanProcessor"),
    _trackCollectionParameters(0), _nGoodFits(0), _nBadFits(0)
{

	_description = "To Be Filled: a Kalman filter track fitter";

	registerInputCollection(     LCIO::TRACK,
				    "InputSeedTracks",
				    "The name of the input collection of track candidates (default: TPCSeedTracks)",
				    _input_seed_tracks_collection_name,
				    std::string("TPCSeedTracks"));

	registerOutputCollection(    LCIO::TRACK,
				    "OutputTracks",
				    "The name of the output collection with the fitted tracks(default: TPCTracks)",
				    _output_tracks_collection_name,
				    std::string("TPCTracks"));

	registerProcessorParameter( "SetOutputTransient" , 
				    "if not 0 the output hits collection is set transient (default: 0)",
				    _outputIsTransient,
				    int(0));

	registerProcessorParameter( "TransDefocussing",
				    "The transverse defocussing in the readout structure sigma_0,trans(in mm) (default: 1)",
				    _transDefocussing,
				    float(1.));
  
	registerProcessorParameter( "TransDiffusionCoef",
				    "The transverse diffusion coefficient (in mm/sqrt(mm) ) (default: 0)",
				    _transDiffusionCoef,
				    float(0.));

	registerProcessorParameter( "LongDefocussing",
				    "The longitudinal defocussing in the readout structure sigma_0,long (in mm) (default: 1)",
				    _longDefocussing,
				    float(1.));

	registerProcessorParameter( "LongDiffusionCoef",
				    "The longitudinal diffusion coefficient (in mm/sqrt(mm) ) (default: 0)",
				    _longDiffusionCoef,
				    float(0.));
}


TrackFitterKalmanProcessor::~TrackFitterKalmanProcessor() 
{
  delete _trackCollectionParameters;
}


void TrackFitterKalmanProcessor::init() {

  printParameters();
  
  // initialise the collection parameters
  _trackCollectionParameters = new IMPL::LCParametersImpl();

  // set the pamameters for this fitter
  _trackCollectionParameters->setValue("TrackFitterType",
				       "KALMAN");

  _trackCollectionParameters->setValue("SimpleChiSquare_TransDefocussing",
				       _transDefocussing);
  _trackCollectionParameters->setValue("SimpleChiSquare_LongDefocussing",
				       _longDefocussing);
  _trackCollectionParameters->setValue("SimpleChiSquare_TransDiffusionCoef",
				       _transDiffusionCoef);
  _trackCollectionParameters->setValue("SimpleChiSquare_LongDiffusionCoef",
				       _longDiffusionCoef);

}


void TrackFitterKalmanProcessor::processRunHeader(LCRunHeader *run) {

    run->parameters().setValue( _processorName + "_revision", "$Rev:$");

    for ( ProcParamMap::iterator i = _map.begin(); i != _map.end(); i++ ) {
        if ( ! i->second->isOptional() || i->second->valueSet() ) {
            run->parameters().setValue( _processorName + "_" + i->second->name(), i->second->value() );
        }
    }

}

void TrackFitterKalmanProcessor::processEvent(LCEvent *evt) {

    streamlog_out(DEBUG4) << "processEvent" << std::endl;
    
    IMPL::LCCollectionVec* seedTracks = 0;
    try {
	// the tracker hit collection
	seedTracks = 
	  dynamic_cast<IMPL::LCCollectionVec*>
	  (evt->getCollection(_input_seed_tracks_collection_name));
    }
    catch (DataNotAvailableException &) {
	streamlog_out(MESSAGE1) << "No seed tracks collection \"" << _input_seed_tracks_collection_name
				<< "\" in event " << evt->getEventNumber() << std::endl;
	return;
    }

    // a collection to hold the new track
    IMPL::LCCollectionVec* outputTrackCollection = new IMPL::LCCollectionVec(EVENT::LCIO::TRACK);

    // set flag that the hits are stored in with the track
    // make sure the hit collection is not transient
    LCFlagImpl trkFlag(0) ;
    trkFlag.setBit( LCIO::TRBIT_HITS ) ;
    outputTrackCollection->setFlag( trkFlag.getFlag() );

    // set parameters for the output collection
    // (rather complicated because the assignment operator cannot be use (it's lcio :-( )
    
    StringVec integerKeys;
    _trackCollectionParameters->getIntKeys(integerKeys);
    for (StringVec::iterator it=integerKeys.begin();
	 it < integerKeys.end(); it ++)
    {
      outputTrackCollection->parameters().setValue( *it,
			      _trackCollectionParameters->getIntVal(*it) );
    }

    StringVec floatKeys;
    _trackCollectionParameters->getFloatKeys(floatKeys);
    for (StringVec::iterator it=floatKeys.begin();
	 it < floatKeys.end(); it ++)
    {
      outputTrackCollection->parameters().setValue( *it,
			      _trackCollectionParameters->getFloatVal(*it) );
    }

    StringVec stringKeys;
    _trackCollectionParameters->getStringKeys(stringKeys);
    for (StringVec::iterator it=stringKeys.begin();
	 it < stringKeys.end(); it ++)
    {
      outputTrackCollection->parameters().setValue( *it,
			      _trackCollectionParameters->getStringVal(*it) );
    }

    streamlog_out(DEBUG3) << "requesting fitter"<<std::endl;

    // request a fitter from the factory
    // a pointer to this event is passed so conditions data can be taken into
    // account
    TrackFitterBase *fitter = TrackFitterFactory::getTrackFitter( TrackFitterBase::FitterTypes::KALMAN,
								 _trackCollectionParameters);

    streamlog_out(DEBUG3) << "fitter is at "<<fitter << " fitter type "
			  << fitter->getFitterType()<<std::endl;

    //loop all tracks
    for (IMPL::LCCollectionVec::iterator trackIter=seedTracks->begin(); 
	 trackIter < seedTracks->end() ; trackIter++)
    {
      EVENT::Track *inputTrack = dynamic_cast<EVENT::Track *>(*trackIter);

      if (inputTrack == 0)
      {
	streamlog_out(WARNING2) << "bad_cast: Track candidates collection"
				<< _input_seed_tracks_collection_name
				<< " does not contain EVENT::Tracks!"<< std::endl;
	throw DataNotAvailableException("No track candidates found.");
      }
	

      streamlog_out(DEBUG3) << "fitting track" << std::endl;

      IMPL::TrackImpl* fittedTrack = fitter->fitTrack(inputTrack);

      // check if fitting was successful
      if ( LCFlagImpl(fittedTrack->getType()).bitSet( TrackFitterBase::FITFAILEDBIT ) )
      {
	_nBadFits++;
      }
      else
      {
	_nGoodFits++;
      }

      streamlog_out(DEBUG3) << "fitting done"<<std::endl;
 
     //determine dEdx:
      //===============
      //FIXME: This currently calculates the mean charge per hit,
      // not the dE/dx because it is difficult to calculate the actual
      // track length in a general formula
      
      TrackerHitVec const & hitVec = fittedTrack->getTrackerHits();
      
      //      double phi =  outputTrack->getPhi();
      //      double d0 =  outputTrack->getD0();
      double Q_sum = 0;
      
      for ( TrackerHitVec::const_iterator hit = hitVec.begin(); 
	    hit < hitVec.end(); hit++ )
      {
	Q_sum += (*hit)->getdEdx();
      }
      
      fittedTrack->setdEdx( Q_sum/hitVec.size() );
      
      outputTrackCollection->push_back(fittedTrack);

    } // loop tracks

    evt->addCollection(outputTrackCollection, _output_tracks_collection_name);

}// processEvent


Processor* TrackFitterKalmanProcessor::newProcessor() {

	return new TrackFitterKalmanProcessor();
}


void TrackFitterKalmanProcessor::check(LCEvent *) {
}


void TrackFitterKalmanProcessor::end() {
  streamlog_out(MESSAGE4) << "Number of good fits: " << _nGoodFits 
			  << "number of bad fits: "  << _nBadFits << std::endl;
#if 1
  TNtupleD *hTracks  = static_cast<TNtupleD *>(gROOT->FindObject("hTracks"));
  TNtupleD *hResXin  = static_cast<TNtupleD *>(gROOT->FindObject("hResXin"));
  TNtupleD *hResZin  = static_cast<TNtupleD *>(gROOT->FindObject("hResZin"));
  hTracks->Write();
  hResXin ->Write();
  hResZin ->Write();
#endif
}

} // namespace marlintpc