//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Add more trackreps to existing tracks
//      Used for adding additional particle hypotheses before
//      final fit
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer E18,   TUM        (original author)
//
//-----------------------------------------------------------

#include "AddTrackRepTask.h"

#include "TClonesArray.h"

#include "GFTrack.h"
#include "GFAbsTrackRep.h"
#include "RKTrackRep.h"


#include <stdio.h>
#include <cmath>

AddTrackRepTask::AddTrackRepTask() 
  :
  fTrackBranchName("meh")
{
  fPdgMap[TString("pion")] = 211;
  fPdgMap[TString("proton")] = 2212;
  fPdgMap[TString("kaon")] = 321;
}

AddTrackRepTask::~AddTrackRepTask() {
  ;
}


InitStatus
AddTrackRepTask::Init()
{
  
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();

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



void
AddTrackRepTask::Exec(Option_t* opt)
{
 
  unsigned int ntracks=fTrackArray->GetEntriesFast();
  
  for(unsigned int it=0; it<ntracks; it++) {
    GFTrack* itr = (GFTrack*) fTrackArray->At(it);
    double sign = itr->getCharge()/std::fabs(itr->getCharge());
    GFAbsTrackRep* rep = itr->getCardinalRep();
    if (dynamic_cast<const RKTrackRep*>(rep) == NULL) {
       Error("AddTrackRepTask::Exec","Wrong track rep; currently only implemented for RKTrackRep!");
    }
    
    //check for trackreps already there:
    unsigned int nReps = itr->getNumReps();
    std::map<int, bool> usedPdgs;
    for(unsigned int nr=0; nr<nReps; nr++) {
      int updg = ((RKTrackRep*)itr->getTrackRep(nr))->getPDG();
      if(updg<0)
	updg = -1*updg;
      usedPdgs[updg] = true;
    }

    
    //add new hypotheses
    for(unsigned int iadd=0; iadd<fToAdd.size(); iadd++) {
      int add = fPdgMap[(fToAdd[iadd])];
      //check if already there     
      if(usedPdgs.count(add) != 0) {
	char buffer[100];
	sprintf(buffer,"PDG %d already present, ignoring wish to add",add);
	Warning("AddTrackRepTask::Exec", TString(buffer));
	continue;
      }
      
      RKTrackRep* newRep = (RKTrackRep*) rep->clone();
      std::cout<<"AddTrackRepTask::Exec() Adding rep with PDG "<<(int)sign*add<<
	std::endl;
      newRep->setPDG((int)sign*add);
      itr->addTrackRep(newRep);
      
    }
    
  }
  
}

void
AddTrackRepTask::AddHypothesis(const TString& hyp) {
  if(fPdgMap.count(hyp) == 0) {
     Error("AddTrackRepTask::AddHypothesis","Hypothesis is inknown in my internal definitions!");
  }
  else {
    fToAdd.push_back(hyp);
  }
}



ClassImp(AddTrackRepTask)