//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class KalmanTask
//      see KalmanTask.hh for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sebastian Neubert    TUM            (original author)
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "KalmanTask.h"
#include <cmath>
// C/C++ Headers ----------------------
#include <algorithm>
#include <iostream>
#include <assert.h>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "TpcCluster.h"
#include "TpcPlanarRecoHit.h"
#include "TpcSPHit.h"

#ifdef WITH_FOPISTUFF
#include "CdcHit.h"
#include "PseudoSpacePoint.h"
#include "CdcSpacepoint.h"
#include "RpcPixHit.h"
#include "BarPixHit.h"
#endif

#include "GFRecoHitFactory.h"
#include "GFKalman.h"
#include "GFDaf.h"
#include "GFException.h"
#include "TH1D.h"
#include "TFile.h"
#include "TGeoTrack.h"
#include "TGeoManager.h"

#include "RecoHits/GFAbsRecoHit.h"
#include "TVector3.h"


#include <signal.h>
#include <stdlib.h>



void sighandler(int sig){
  std::cerr << "sighandler for sig " << sig << std::endl;
  abort();
}


// Class Member definitions -----------

 
KalmanTask::KalmanTask()
  : FairTask("Kalman Filter"), _persistence(kFALSE),_lazy(0),_numIt(1), 
    _trackBranchName("TrackPreFit"), _outBranchName("TrackPostFit"), 
    fMVDPixelBranchName("MVDHitsPixel"), fMVDStripBranchName("MVDHitsStrip"),
    fFopiCDCBranchName("CdcHit"), fTpcClusterBranchName("TpcCluster"),
    fRpcPixHitBranchName("RpcPixHit"),
    fBarPixHitBranchName("BarPixHit"),
    fUseDAF(false)
    
{
  fVerbose = 0;
}

 
KalmanTask::~KalmanTask()
{
  if(_pH!=NULL)delete _pH;
  if(_chi2H!=NULL)delete _chi2H;
  delete fFitter;
}

InitStatus
KalmanTask::Init()
{
  
  // Fitting ---------------- can go to another task!
  
  if(!fUseDAF)  {
    fFitter = new GFKalman();
    ((GFKalman*)fFitter)->setNumIterations(_numIt);
  }
  else
    fFitter = new GFDaf();
  
  for(unsigned int n=0; n<fIgnorePDGs.size(); n++) 
    fFitter->addIgnorePDG(fIgnorePDGs[n]);
  
    
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();

  if(ioman==0) {
    Error("KalmanTask::Init","RootManager not instantiated!");
    return kERROR;
  }
  
  // Get input collection
  _trackArray=(TClonesArray*) ioman->GetObject(_trackBranchName);
  if(_trackArray==0) {
    Error("KalmanTask::Init","track-array not found!");
    return kERROR;
  }
  

  // Build hit factories -----------------------------
  _theRecoHitFactory = new GFRecoHitFactory();
  _clusterArray = (TClonesArray*) ioman->GetObject(fTpcClusterBranchName);
   if(_clusterArray==0){
     Error("KalmanTask::Init","TpcCluster array not found");
     return kERROR;
   }
   else{ 
     TString type(_clusterArray->GetClass()->GetName());
     if(type == "TpcCluster")
       _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fTpcClusterBranchName),
				       new GFRecoHitProducer<TpcCluster,TpcSPHit>(_clusterArray));
     else if(type == "TpcSPHit")
       _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fTpcClusterBranchName),
				       new GFRecoHitProducer<TpcSPHit,TpcSPHit>(_clusterArray));
     else {
       TString err("Invalid Object type in TPC hit array: ");
       err += type;
       Fatal("KalmanTask::Init",err);
     }
   }
   
   #ifdef WITH_FOPISTUFF
   TClonesArray* cdcHits = (TClonesArray*) ioman->GetObject(fFopiCDCBranchName);
   if(cdcHits==NULL){
     Error("KalmanTask::Init","CDCHit array not found");
   } else {
     TString type(cdcHits->GetClass()->GetName());
     bool gotId = false;
     if(type == "CdcHit") { 
       _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fFopiCDCBranchName),
				       new GFRecoHitProducer<CdcHit,PseudoSpacePoint>(cdcHits));
       gotId = true;
     }
     if(type == "PseudoSpacePoint") {
       _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fFopiCDCBranchName),
				       new GFRecoHitProducer<PseudoSpacePoint,PseudoSpacePoint>(cdcHits));
       gotId = true;
     }
     if(type == "CdcSpacepoint") {
       _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fFopiCDCBranchName),
				       new GFRecoHitProducer<CdcSpacepoint,CdcSpacepoint>(cdcHits));
       gotId = true;
     }
     if(!gotId) {
       TString err("Invalid Object type in CDC hit array: ");
       err += type;
       Fatal("KalmanTask::Init",err);
     }
     
   }

   _rpcHitArray = (TClonesArray*) ioman->GetObject(fRpcPixHitBranchName);
   if(_rpcHitArray==NULL) {
     Error("KalmanTask::Init","RpcPixHit array not found");
   }else{
	   _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fRpcPixHitBranchName),
				   	   new GFRecoHitProducer<RpcPixHit,RpcPixHit>(_rpcHitArray));
   }

   _barHitArray = (TClonesArray*) ioman->GetObject(fBarPixHitBranchName);
   if(_barHitArray==NULL) {
     Error("KalmanTask::Init","BarPixHit array not found");
   }else{
	   _theRecoHitFactory->addProducer(FairRootManager::Instance()->GetBranchId(fBarPixHitBranchName),
			   new GFRecoHitProducer<BarPixHit,BarPixHit>(_barHitArray));
   }
   #endif
   
     
   _trackOutArray = new TClonesArray("GFTrack");
   ioman->Register(_outBranchName,"", _trackOutArray, _persistence);


  // setup histograms
  _pH=new TH1D("pH","p",100,0.4,0.6);
  _chi2H=new TH1D("chi2H","chi2",100,0,20);
  _massV0=new TH1D("massV0","massV0",100,0,5);
  _massETAC=new TH1D("massEta","massEta",100,2.5,3.5);

  std::cout<<"Kalman Task Init success\n";
  return kSUCCESS;
}


void
KalmanTask::Exec(Option_t* opt)
{
  if (fVerbose) std::cout<<"KalmanTask::Exec"<<std::endl;

  // Reset output Array
  _trackOutArray->Delete();
  if(_trackArray==0) Fatal("Kalman::Exec)","No TrackArray");


  Int_t ntracks=_trackArray->GetEntriesFast();
  
  if(ntracks>20000){
    std::cout<<"ntracks="<<ntracks<<" Evil Event! skipping"<<std::endl;
    return;
  }

  
  
  //  fFitter.setLazy(_lazy);
  


  for(Int_t itr=0;itr<ntracks;++itr){
    if (fVerbose) std::cout<<"starting track"<<itr<<std::endl;
    GFTrack* trk=(GFTrack*)_trackArray->At(itr);
    
    // Load RecoHits 
    try {
      trk->addHitVector(_theRecoHitFactory->createMany(trk->getCand()));
      if (fVerbose) std::cout<<trk->getNumHits()<<" hits in track " <<itr<<std::endl;
    }
    catch(GFException& e) {
      if(fVerbose) {
	std::cout << e.what() << std::endl;
	e.info();
      }
      throw e;
    }
    if (fVerbose) std::cout<<"*** Number of clusters in track: "<<trk->getNumHits()<<" ***"<<std::endl;

    Int_t size = _trackOutArray->GetEntriesFast();
    GFTrack* trkCopy = new((*_trackOutArray)[size]) GFTrack(*trk);

    
    // Start Fitter
    try{
      if (fVerbose) std::cerr << "KalmanTask::Exec - Calling processTrack" << std::endl;
      fFitter->processTrack(trkCopy);
    }
    catch (GFException& e){
      if(fVerbose)
	std::cout<<e.what()<<std::endl;
    }
    

    // Print Track Parameters after fit
    if(trkCopy->getTrackRep(0)->getStatusFlag()==0){
      if(fVerbose) trkCopy->getTrackRep(0)->Print();
      double p=trkCopy->getMom().Mag();
      _pH->Fill(p);
      
      double chi2=trkCopy->getChiSqu();
      _chi2H->Fill(chi2);
    }

  }

  if (fVerbose) {
    std::cout<<"Fitting done"<<std::endl;
    std::cout<<"Fitted "<< _trackOutArray->GetEntriesFast() << " tracks" << std::endl;
  }
  
  return;
}


void 
KalmanTask::WriteHistograms(const TString& filename){
  TFile* file = new TFile(filename,"UPDATE");
  file->mkdir("Kalman");
  file->cd("Kalman");

  _pH->Write();
  delete _pH;
  _pH=NULL;

  _chi2H->Write();
  delete _chi2H;
  _chi2H=NULL;

  _massV0->Write();
  delete _massV0;
  _massV0=NULL;

  _massETAC->Write();
  delete _massETAC;
  _massETAC=NULL;

  file->Close();
  delete file;
}


ClassImp(KalmanTask)