#include "PndTorinoDetectorKalmanTask.h"

// torino detector
#include "PndTorinoDetectorHit.h"
#include "PndTorinoDetectorRecoHit.h"
// track in pandaroot
#include "PndTrack.h"
#include "PndTrackID.h"
#include "PndTrackCand.h"
// genfit
#include "GFKalman.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "GFRecoHitProducer.h"
// adapters
#include "PndGenfitAdapters.h"
// track rep
#include "GeaneTrackRep.h"


using namespace std;

PndTorinoDetectorKalmanTask::PndTorinoDetectorKalmanTask() : FairTask(), fTrackArray(NULL), fTrackIDArray(NULL), fHitArray(NULL), fFitTrackArray(new TClonesArray("PndTrack")), fPDGcode(13) {}

PndTorinoDetectorKalmanTask::~PndTorinoDetectorKalmanTask() {}

InitStatus PndTorinoDetectorKalmanTask::Init() {

  fEventCounter = 0;

  // instance of FairRootManager
  FairRootManager  *ioman = FairRootManager::Instance();
  if (!ioman) {
    cout << "-E- RootManager not instantised!" << endl;
    return kFATAL;
  }

  // get hit array
  fHitArray  = (TClonesArray*) ioman->GetObject("PndTorinoDetectorHit");
  if (!fHitArray) {
    cout << "-E- PndTorinoDetectorKalman::Init: No hit array!" << endl;
    return kERROR;
  }

  // get track array
  fTrackArray  = (TClonesArray*) ioman->GetObject("PndTorinoDetectorTrack");
  if (!fTrackArray) {
    cout << "-E- PndTorinoDetectorKalman::Init: No track array!" << endl;
    return kERROR;
  }

  // get trackID array
  fTrackIDArray  = (TClonesArray*) ioman->GetObject("PndTorinoDetectorTrackID");
  if (!fTrackIDArray) {
    cout << "-E- PndTorinoDetectorKalman::Init: No trackID array!" << endl;
    return kERROR;
  }

  // ---------------- reco hit factory
  fTheRecoHitFactory = new GFRecoHitFactory();
  // recoFactory->addProducer(detectorID, 
  //                          recohitproducer<hit, recohit> (TCA))
  fTheRecoHitFactory->addProducer(ioman->GetBranchId("PndTorinoDetectorHit"), 
				  new GFRecoHitProducer<PndTorinoDetectorHit, PndTorinoDetectorRecoHit>(fHitArray));

 
  //  if(fVerbose == 0) GFException::quiet(true);
   // ---------------- geane
  fPro = new FairGeanePro();


  // register output -----------------------------------------
  ioman->Register("PndTorinoDetectorKalmanTrack", "PndTorinoDetector", fFitTrackArray, kTRUE);

  return kSUCCESS;

}


void PndTorinoDetectorKalmanTask::Exec(Option_t* option) {

  if(fVerbose > 0)  cout << "####### Kalman on event " << fEventCounter << endl;
  fEventCounter++;
  fFitTrackArray->Delete(); 

  // prepare output track
  PndTrack *fittrack = NULL;

  // loop over found tracks
  for(int itrk = 0; itrk < fTrackArray->GetEntriesFast(); itrk++) {
  if(fVerbose > 0)  cout << "------- Kalman on track " << itrk << endl;

    // ============================================================================
    /** we start from the track parameters we found during the track finding **/
    PndTrack *track = (PndTrack*) fTrackArray->At(itrk);
    if(!track) continue;

    /** to fit only primaries
	PndTrackID *trackID = (PndTrackID*) fTrackIDArray->At(itrk);
	if(!trackID) continue;
	if(trackID->GetCorrTrackID() > 2) continue;
    **/
    
    // get params at 1st hit
    FairTrackParP par = track->GetParamFirst();
    
    Int_t  charge = par.GetQ();
    // correct pdgcode according to charge
    SetRealPdg(fPDGcode, charge);

    // get initial values from prefit
    TVector3 StartPos = par.GetPosition();
    TVector3 StartMom = par.GetMomentum();
    TVector3 StartPosErr(par.GetDX(), par.GetDY(), par.GetDZ()); 
    TVector3 StartMomErr(par.GetDPx(), par.GetDPy(), par.GetDPz());


    if(fVerbose) {
      cout << "trackcand has " << track->GetTrackCand().GetNHits() << " hits" << endl;
      cout << "startpos" << endl;
      StartPos.Print();
      cout << "startmom" << endl;
      StartMom.Print();
    }

    // ============================================================================
    /** we want to start from the POCA to origin --> extrapolation back **/

    // set point = IP
    fPro->SetPoint(TVector3(0., 0., 0.));
    fPro->PropagateToPCA(1, -1);

    // transform the ParP to ParH
    Int_t ierr = 0;
    FairTrackParH *helix = new FairTrackParH(&par, ierr);

    // set final par container
    FairTrackParH *atvertex = new FairTrackParH();

    Bool_t rc =  fPro->Propagate(helix, atvertex, fPDGcode);
    if(!rc) {
      cout << "PROPAGATION FAILED " << ierr << endl;
      helix->GetPosition().Print();
      helix->GetMomentum().Print();
    }
    else     cout << "PROPAGATION SUCCESS " << ierr << endl;

    // ============================================================================
    // if the propagation is successful
    if (rc)
      {
	StartPos.SetXYZ(atvertex->GetX(), atvertex->GetY(), atvertex->GetZ());
	StartMom.SetXYZ(atvertex->GetPx(), atvertex->GetPy(), atvertex->GetPz());
	StartPosErr.SetXYZ(atvertex->GetDX(), atvertex->GetDY(), atvertex->GetDZ());
	StartMomErr.SetXYZ(atvertex->GetDPx(), atvertex->GetDPy(), atvertex->GetDPz());
      }


    // ===========================================================================
    // ========================== KALMAN FIT =====================================

    // define starting plane
    GFDetPlane startplane(StartPos, TVector3(1.,0.,0.), TVector3(0.,1.,0.));

    // define track rep
    GeaneTrackRep *grep = new GeaneTrackRep(fPro,
					    startplane, StartMom,
					    StartPosErr, StartMomErr,
					    charge, fPDGcode);
    grep->setPropDir(1);

    // define GFTrackCand
    PndTrackCand trackCand = track->GetTrackCand();
    GFTrackCand *candidate = PndTrackCand2GenfitTrackCand(&trackCand);

    // define GFTrack, with track cand, trackrep & recohits
    GFTrack* trk= new GFTrack(grep); // <-- trackrep
    trk->setCandidate(*candidate);    // <-- track cand
   
    try {
      trk->addHitVector(fTheRecoHitFactory->createMany(*candidate)); // <-- recohits
    }
    catch(GFException& e)
      {
	cout << "-E- PndTorinoDetectorKalmanTask::Exec, genfit exception: trk->addHitVector " << endl;
	cout << e.what() << endl;
      }
    
    // KALMAN <<===================================================================
    try {
      GFKalman genfitter;
      genfitter.processTrack(trk);
    }
    catch (GFException e)
      {
	cout << "-E- PndTorinoDetectorKalmanTask::Exec, genfit exception: processTrack " << e.what() << endl;
	cout << e.what() << endl;
  	fittrack = track;
	fittrack->SetFlag(-10); // flag -10: kalman failed
      }

    // ============================================================================
    // conversion from GFTrack to PndTrack
    try   { 
      fittrack = (PndTrack*)GenfitTrack2PndTrack(trk);
    }
    catch (GFException e)
      {
	cout << "-E- PndTorinoDetectorKalmanTask::Exec, PndGenfitAdapters exception: conversion filed " << endl;
	cout<<e.what()<<endl;
	fittrack = track;
	fittrack->SetFlag(-2); // flag -2: conversion failed
      } 

    // ============================================================================
    // register to output
    Int_t size = fFitTrackArray->GetEntriesFast();
    fittrack = new((*fFitTrackArray)[size]) PndTrack(fittrack->GetParamFirst(), fittrack->GetParamLast(), 
						     fittrack->GetTrackCand(),
						     fittrack->GetFlag(), fittrack->GetChi2(), fittrack->GetNDF(), 
						     fittrack->GetPidHypo(), 
						     itrk, FairRootManager::Instance()->GetBranchId("PndTorinoDetectorTrack"));
  }

 
  return;

}
  

void PndTorinoDetectorKalmanTask::SetRealPdg(Int_t pdg, Int_t charge)
{  
  switch (abs(pdg))
    {
    case 11:
      fPDGcode = charge * (-11);
      break;
    case 13:
      fPDGcode = charge * (-13);
      break;
    case 211:
      fPDGcode = charge * (211);
      break;
    case 321:
      fPDGcode = charge * (321);
      break;
    case 2212:
      fPDGcode = charge * (2212);
      break;
    default:
      cout << "-I- PndTorinoDetectorKalmanTask::SetParticleHypo: Not recognised PID set -> Using default MUON hypothesis" << endl;
      fPDGcode = charge * (-13);
      break;
    }
}
ClassImp(PndTorinoDetectorKalmanTask);