//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcTrackInitTask
//      Construct TpcSP RecoHits and GFTrack objects 
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sebastian Neubert    TUM            (original author)
//      Johannes Rauch       TUM
//      Felix Boehmer        TUM
//
//-----------------------------------------------------------

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

// This Class' Header ------------------
#include "TpcTrackInitTask.h"

// C/C++ Headers ----------------------
#include <map>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "FairRuntimeDb.h"
#include "TClonesArray.h"
#include "TpcCluster.h"
#include "TpcSPHit.h"
#include "TpcRiemannHit.h"
#include "TpcRiemannTrack.h"
#include "McIdCollection.h"
#include "TpcAlignmentManager.h"

#include "TpcDigiPar.h"

#include "GFTrackCand.h"
#include "GFTrack.h"
#include "RKTrackRep.h"

#include "TVector3.h"
#include "FairMCPoint.h"
#include "FairRunAna.h"
#include "GFDetPlane.h"
#include "TDatabasePDG.h"
#include "FairField.h"
#include "PndConstField.h"
#include "PndMultiField.h"
#include "PndFieldAdaptor.h"
#include "GFFieldManager.h"
#include "PndMCTrack.h"

#include "TFile.h"
#include "TVector3.h"
#include "TMath.h"


#include <cmath>


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

ClassImp(TpcTrackInitTask)

TpcTrackInitTask::TpcTrackInitTask()
: FairTask("Tpc Track Init Task"),
  _persistence(kFALSE),
  _arhpersistence(kFALSE),

  _clusterBranchName("TpcCluster"),
  _riemannBranchName("RiemannTrack"),
  _trackBranchName("TrackPreFit"),
  _recoHitBranchName("TpcSPHit"),

  _cosmicSorting(false),
  _smoothing(true),
//_geane(false),
  _useFullCov(false),
  _weightedPlaneConstruction(false),
  _cutCov(false),
  _mcPid(false),
  _pdg(211),    //default
  _issim(false),
  _dipCut(0.01),

  counter(0),
  Bz(0),

  fmaxLenZbackw(0),
  
  fClBrID(-1),
  fSPHitID(-1),
  fAlign(kTRUE),
  fInteractionPoint(0,0,42.78)
  {
    ;
  }

TpcTrackInitTask::~TpcTrackInitTask()
{
  ;
}


InitStatus
TpcTrackInitTask::Init()
{
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();
  if(ioman==0){
    Error("TpcTrackInitTask::Init","RootManager not instantiated!");
    return kERROR;
  }

  // get input arrays
  if (_mcPid){
    _mcTrackArray=(TClonesArray*) ioman->GetObject("MCTrack");
    if(_mcTrackArray==0){
      Warning("TpcTrackInitTask::Init","MCTrack-array not found! Cannot use ideal PID");
      _mcPid=false;
    }
  }

  _clusterArray=(TClonesArray*) ioman->GetObject(_clusterBranchName);
  if(_clusterArray==0){
    Fatal("TpcTrackInitTask::Init","Cluster-array not found!");
    return kERROR;
  }
  // get TPC "detector ID". In fact, it is the branch ID of the TPC Clusters.
  fTpcClusterBranchID = FairRootManager::Instance()->GetBranchId(_clusterBranchName);


  _riemannTrackArray=(TClonesArray*) ioman->GetObject(_riemannBranchName);
  if(_riemannTrackArray==0){
    Fatal("TpcTrackInitTask::Init","RiemannTrack-array not found!");
    return kERROR;
  }

  // create and register output arrays
  _trackArray = new TClonesArray("GFTrack");
  ioman->Register(_trackBranchName,"GenFit",_trackArray,_persistence);

  _recoHitArray = new TClonesArray("TpcSPHit");
  ioman->Register(_recoHitBranchName,"tpc",_recoHitArray,_arhpersistence || _persistence);

  fSPHitID = FairRootManager::Instance()->GetBranchId(_recoHitBranchName);
  fClBrID = FairRootManager::Instance()->GetBranchId(_clusterBranchName);
  
  //get the magnetic field for curvature seeding
  FairField* field=FairRunAna::Instance()->GetField();
  bool CField = dynamic_cast<PndConstField*>(field);
  bool MField = dynamic_cast<PndMultiField*>(field);
  //GFFieldManager::getInstance()->init(new PndFieldAdaptor(field));
  if(MField) {
    Double_t O[3], B[3];
    O[0]=0; O[1]=0; O[2]=0;
    field->GetFieldValue(O,B);
    Bz=B[2];
    std::cerr<<"TpcTrackInitTask: "<<"No const field! Curvature seeding not valid... Setting Bz = "<<Bz<<std::endl;
  }
  else if(CField) {
    Bz=field->GetBz(0.,0.,0.);
    std::cerr<<"TpcTrackInitTask: "<<"const field! Setting Bz = "<<Bz<<std::endl;
  }
  else{
    Bz=20.;
    std::cerr<<"TpcTrackInitTask: "<<"default setting Bz = "<<Bz<<std::endl;
  }

  //init FairGeanePro
  //if(_geane) gPro = new FairGeanePro();

  return kSUCCESS;
}


void
    TpcTrackInitTask::SetParContainers() {

  if (fVerbose) {
    std::cout<<"TpcClusterFinderTask::SetParContainers"<<std::endl;
    std::cout.flush();
  }

  // Get run and runtime database
  FairRun* run = FairRun::Instance();
  if ( ! run ) Fatal("SetParContainers", "No analysis run");

  FairRuntimeDb* db = run->GetRuntimeDb();
  if ( ! db ) Fatal("SetParContainers", "No runtime database");

  // Get Tpc digitisation parameter container
  fpar= (TpcDigiPar*) db->getContainer("TpcDigiPar");
  if (! fpar ) Fatal("SetParContainers", "TpcDigiPar not found");
}



void
    TpcTrackInitTask::Exec(Option_t* opt)
{
  if (fVerbose)
    std::cout << "TpcTrackInitTask::Exec; Event Number: " << counter++ << std::endl;
  
  // Reset output Arrays
  if(_trackArray==0) Fatal("TpcTrackInitTask::Exec)","No TrackOutArray");
  _trackArray->Delete();
  if(_recoHitArray==0) Fatal("TpcTrackInitTask::Exec)","No RecoHitOutArray");
  _recoHitArray->Delete();

  fRecoHitMap.clear();
  fToWrite.clear();

  //First, loop over all clusters and build recoHits
  unsigned int nCl = _clusterArray->GetEntriesFast();
  for(unsigned int iCl=0; iCl<nCl; iCl++) {
    TpcCluster* cl = (TpcCluster*) _clusterArray->At(iCl);
    TpcSPHit* theHit = new TpcSPHit(cl,_weightedPlaneConstruction,_cutCov);
    //if (_useFullCov)
    //theHit->setHitCov(cl->cov());
    theHit->setClusterBranchID(fClBrID);
    fRecoHitMap[cl] = theHit;
  }

  // build GFTrackCands
  std::vector<GFTrackCand*> candlist;

  double pbackup = 2.;  // momentum value that is set when other initialisations fail

  // loop over Riemann tracks
  unsigned int nr = _riemannTrackArray->GetEntriesFast();
  if (fVerbose)
    std::cout << "Looping over " << nr << " riemann tracks" << std::endl;

  for(unsigned int itrk=0; itrk<nr; ++itrk){
    TpcRiemannTrack* trk = (TpcRiemannTrack*)_riemannTrackArray->At(itrk);
    
    int nhits=trk->getNumHits();
    if (fVerbose) std::cout<<"Tracklet "<<itrk;

    // check if track is pre-fitted
    if (! trk->isFitted()) {
      if (fVerbose) std::cout << " - skipping, track not prefitted" << std::endl;
      continue;
    }
    // check if track too steep
    double trackSinDip = trk->sinDip();
    double p = trk->getMom(Bz);
    double trackR = trk->r();

    if (_issim)
      {	
	if (TMath::Abs(trackSinDip)<_dipCut or trackR<0.5) 
	  {
	    if (fVerbose) std::cout << " sin(dip) too small->Set p backup: " << trackSinDip << std::endl;
	    p=pbackup;
	  }
      }
    else
      {
	// check if track too steep
	if (TMath::Abs(trackSinDip)<_dipCut) {
	  if (fVerbose) std::cout << " - skipping, sin(dip) too small: " << trackSinDip << std::endl;
	  continue;
	}
	// ceck if momentum high enough
	if (Bz==0) p=pbackup;
	if(p<5.E-2) {  // 50 MeV
	  if (fVerbose) std::cout << " - skipping, momentum too small: " << p*1E3 << " MeV" << std::endl;
	  continue;
	}
      }
    unsigned int trackId(trk->mcid().DominantID().mctrackID());
    int eventId(trk->mcid().DominantID().mceventID());
    
    //if(eventId!=0)trackId+=10000;

    // check pdg
    int pdg(_pdg);
    if(_mcPid) pdg=((PndMCTrack*)(_mcTrackArray->At(trackId)))->GetPdgCode();
    if (fVerbose)
    std::cerr<<"##############################  TrackInit for PDGID: "<<pdg<<"   ##############################"<<std::endl;

    double pdgCharge(TDatabasePDG::Instance()->GetParticle(pdg)->Charge()/3.);

    int winding(trk->winding()); // we look in z direction!

    if (pdgCharge < 0) {
      pdg *= -1;
      pdgCharge *= -1.;
    }
    if (winding > 0) {
      pdg *= -1;
      pdgCharge *= -1.;
    }
    if (Bz < 0) {
      pdg *= -1;
      pdgCharge *= -1.;
    }
    
    // check sorting
    bool invertTrack(true);

    TVector3 ps1, ps2, dir1, dir2,dir3;
    Double_t angle;

    trk->getPosDirOnHelix(0, ps1, dir1);
    trk->getPosDirOnHelix(nhits-1, ps2, dir2);
    dir2*=-1.; // point in direction of track
    dir3=(TVector3)(ps1-ps2);
    angle=TMath::RadToDeg()*abs(dir3.Phi());

    // this task runs for physics only or after the event deconv task, where the physics event is at t = 0
    // so we sort the track in a way that they fly away from the IP (will go wrong for some secondaries!)
    // todo: this has to be refined at some point
    if (_cosmicSorting)
    {
        if(ps1.Y()>ps2.Y()) invertTrack = false;
        else invertTrack = true;

	if(TMath::Abs(trackSinDip)<_dipCut or trackR<0.5) 
	  {
	    if(ps1.Z()<ps2.Z()) invertTrack=false;
	    else invertTrack=true;
	  }
      
    }else{
      if ((ps1-fInteractionPoint).Mag() < (ps2-fInteractionPoint).Mag()) invertTrack = false;
      else  invertTrack = true;
    }

    if (invertTrack){
      winding*=-1;
      pdg *= -1;
      pdgCharge *= -1;
    }

    TParticlePDG * part = TDatabasePDG::Instance()->GetParticle(pdg);
    if(part == 0){
      if (fVerbose) std::cout << " - skipping, unknown PDG id: " << pdg;
      continue;
    }

    if (fVerbose) std::cout<<std::endl;

    // create GFTrackCands
    GFTrackCand* cand=new GFTrackCand();

    // fill hits into GFTrackCands  and get seed values
    TVector3 pos1, direction;


    //tracks cannot share hits (see tracking task), so no worries 
    if(!invertTrack){
      for(unsigned int ih=0; ih<nhits; ++ih){
        TpcSPHit* theHit = fRecoHitMap[trk->getHit(ih)->cluster()];
        unsigned int nOut = _recoHitArray->GetEntriesFast();
        new ((*_recoHitArray)[nOut]) TpcSPHit(*theHit);
        fToWrite.insert(theHit);
        cand->addHit(fSPHitID, nOut);
      }
      trk->getPosDirOnHelix(0, pos1, direction);
    }
    else { // invert track
      for(unsigned int ih=nhits; ih>0; --ih){
        TpcSPHit* theHit = fRecoHitMap[trk->getHit(ih-1)->cluster()];
        unsigned int nOut = _recoHitArray->GetEntriesFast();
        new ((*_recoHitArray)[nOut]) TpcSPHit(*theHit);
        fToWrite.insert(theHit);
        cand->addHit(fSPHitID, nOut);
      }
      trk->getPosDirOnHelix(trk->getNumHits()-1, pos1, direction);
      direction *= -1.;
    }// finished filling hits

    TVector3 poserr(1,1,1);
    poserr*=trk->resolution();

    TVector3 mom(p * direction);
    if (_issim)
      {	
	if (TMath::Abs(trackSinDip)<_dipCut or trackR < 0.5) 
	  {
	    mom.SetZ(p);
	  }
      }
    

    TVector3 momerr(fabs(mom.X()),fabs(mom.Y()),fabs(mom.Z()));
    momerr *= trk->resolution()*5.;

   

    if (fVerbose) {
      double trackDip = trk->dip();
      std::cout<<" center of track "; trk->center().Print();
      std::cout<<" Radius of track [cm]: " << trackR;
      std::cout<<"\n Dip of track [deg]:   " << trackDip/TMath::Pi()*180;
      std::cout<<"\n seed values: ";
      std::cout<<"\n  start position: "; pos1.Print();
      std::cout<<"  momentum [GeV]: "<<p;
      std::cout<<"\n  p_perp [GeV]:   " << trackR*0.0003*TMath::Abs(Bz);
      std::cout<<"\n  direction: "; direction.Print();
      std::cout<<"  winding: "<<winding;
      std::cout<<"\n  invertTrack: "<<invertTrack;
      std::cout<<"\n  pdg id: "<<pdg<<std::endl;
    }

    // set seed values to cands
    //    cand->setCurv(1./trackR);
    //    cand->setDip(trk->dip());
    
    //alignment of initial track parameters
    if (fAlign)
    {
      TpcAlignmentManager* alMan=TpcAlignmentManager::getInstance();
      pos1=alMan->localToMaster("tpc",pos1);
      mom=alMan->localToMasterVect("tpc",mom);

      poserr=alMan->localToMasterVect("tpc",poserr);
      momerr=alMan->localToMasterVect("tpc",momerr);
    }

    //cand->setComplTrackSeed(pos1, mom, pdg, poserr, momerr*(1./p));
    cand->setMcTrackId(trackId);

    candlist.push_back(cand);


    //RK TRACKREP
    RKTrackRep* rkrep = new RKTrackRep(pos1, mom, poserr, momerr, pdg);

    // store GFTracks in output array
    GFTrack* gftrk=new((*_trackArray)[_trackArray->GetEntriesFast()]) GFTrack(rkrep);
    gftrk->setCandidate(*cand); // here the candidate is copied!

    //GEANE TRACKREP
    // if(_geane) {
    //   TVector3 u=mom.Orthogonal();
    //   u.SetMag(1.);
    //   TVector3 v=mom.Cross(u);
    //   v.SetMag(1.);
    //   GFDetPlane pl(pos1,u,v);

    //   GeaneTrackRep* grep = new GeaneTrackRep(gPro,pl,mom,poserr,momerr,pdgCharge,pdg);
    //   // add rep //and set as cardinal rep
    //   gftrk->addTrackRep(grep);
    //   //gftrk->setCardinalRep(gftrk->getNumReps()-1);
    // }
    
    //SMOOTHING
    if(_smoothing) gftrk->setSmoothing(true);
    
  }// end loop over tracks
  
  //write out all other recoHits if required, clean up:
  std::map<TpcCluster*, TpcSPHit*>::const_iterator it;
  for(it=fRecoHitMap.begin(); it!=fRecoHitMap.end(); it++) {
    TpcSPHit* iHit = it->second;
    if(!fToWrite.count(iHit)) {
      if(_arhpersistence ) 
        new ((*_recoHitArray)[_recoHitArray->GetEntriesFast()]) TpcSPHit(*iHit);
    }
    delete iHit;  //all of them have been copied during TCA write ...
  }
  
  if (fVerbose) std::cout<<"TpcTrackInitTask::Exec:: "
      <<candlist.size()<<" tracks setup."<<std::endl;

}


void
    TpcTrackInitTask::WriteHistograms(const TString& filename) {
  std::cerr<< "TpcTrackInitTask::WriteHistograms"<<std::endl;

  TFile* file=FairRootManager::Instance()->GetOutFile();
  file->mkdir("TrackInit");
  file->cd("TrackInit");
}