//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Conversion of CdcTracks into GFtrack objects
//      plus initialization
//
// Environment:
//      Software developed for the GEM-TPC in FOPI
//
// Author List:
//      Francesco Cusanno    TUM            (original author)
//      Felix Boehmer, E18   TUM            (Overhaul)
//
//-----------------------------------------------------------

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

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

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

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "CdcHit.h"
#include "CdcTrack.h"
#include "CdcSpacepoint.h"
#include "PseudoSpacePoint.h"

#include "GFTrackCand.h"
#include "GFTrack.h"
#include "RKTrackRep.h"
#include "GFException.h"
#include "TVector3.h"
#include "FairRunAna.h"

#include "TDatabasePDG.h"
#include "PndConstField.h"
#include "PndMultiField.h"
#include "GFFieldManager.h"


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



FopiTrackInitTask::FopiTrackInitTask()
    : fHitBranchName("CdcHit"), 
      fTrackBranchName("CdcTrack"),
      fOutBranchName("CdcTrackPreFit"),
      fRHOutBranchName("CdcSpacepoint"),
      fHitArray(NULL),
      fTrackArray(NULL),
      fOutArray(NULL), 
      fRHOutArray(NULL),
      fPersistence(kTRUE),
      fArhPersistence(kFALSE),
      fTrackPersistence(kFALSE),
      fSmoothing(true),
      fVerbose(false),
      fPDG(false),
      fUsePseudoSP(kFALSE),
      _weightedPlaneConstruction(kFALSE),
      _cutCov(kFALSE),
      fCounter(0),
      fCdcBrId(-1),  
      fPspBrId(-1),
      fMinCDCHits(5),
      fEx(0),fEy(0),fEz(0),
      fConstantErrors(kFALSE),
      fPiKaonThresh(0.5),
      fKaonProtThresh(0.7),
      fExtrapToZ(kFALSE)
{;}   

FopiTrackInitTask::~FopiTrackInitTask()
{;}


InitStatus
FopiTrackInitTask::Init()
{
    //Get ROOT Manager
    FairRootManager* ioman= FairRootManager::Instance();
    if(ioman==NULL){
        Error("FopiTrackInitTask::Init","RootManager not instantiated!");
        return kERROR;
    }
    
    fCdcBrId = FairRootManager::Instance()->GetBranchId(fHitBranchName);
    
    // get input arrays    
    fHitArray=(TClonesArray*) ioman->GetObject(fHitBranchName);
    if(fHitArray==NULL){
        Error("FopiTrackInitTask::Init","CDCHit-Array not found!");
        return kERROR;
    }
    
    fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);
    if(fTrackArray==NULL){
        Error("FopiTrackInitTask::Init","CDCTrack-Array not found!");
        return kERROR;
    }
    
    
    
    // create and register output array
    fOutArray = new TClonesArray("GFTrack");
    ioman->Register(fOutBranchName,"GenFit",fOutArray,fTrackPersistence);
    if(fUsePseudoSP)
        fRHOutArray = new TClonesArray("PseudoSpacePoint");
    else 
        fRHOutArray = new TClonesArray("CdcSpacepoint");
    ioman->Register(fRHOutBranchName,"GenFit",fRHOutArray,fPersistence);
    
    fPspBrId = FairRootManager::Instance()->GetBranchId(fRHOutBranchName);
    
    fCounter=0;
    return kSUCCESS;
}


//helper functor
bool 
sortingbyrad(const CdcHit* lh, const CdcHit* rh) {
    return(lh->GetHitPos().Perp() < rh->GetHitPos().Perp());
}

void
FopiTrackInitTask::Exec(Option_t* opt)
{
    if(fVerbose)
        std::cout<<"FopiTrackInitTask::Exec(): Event Number: "<<fCounter<<"\n";
    fCounter++;
    
    // Reset output Arrays
    if(fOutArray==NULL) Fatal("FopiTrackInitTask::Exec()","No OutTrackArray");
    fOutArray->Delete();
    if(fRHOutArray==NULL) Fatal("FopiTrackInitTask::Exec()","No RHOutTrackArray");
    fRHOutArray->Delete();
    
    if(fVerbose) std::cout<<"FopiTrackInitTask::Exec() ... \n";
    
    fRecoHitMap.clear();
    fToWrite.clear();
    
    
    //create recohits from CdcHits:
    int nCdcHits = fHitArray->GetEntriesFast();
    for(int icdch=0; icdch<nCdcHits; icdch++) {
        CdcHit* ih = (CdcHit*) (*fHitArray)[icdch];
        //unsigned int nOut = fRHOutArray->GetEntriesFast();
        //CdcSpacepoint* ipsp = new ((*fRHOutArray)[nOut]) CdcSpacepoint(ih);
        if(fUsePseudoSP) {
            fRecoHitMap[ih] = new PseudoSpacePoint(ih);
            PseudoSpacePoint* meh = (PseudoSpacePoint*) fRecoHitMap[ih];
            meh->setHitBranchID(fPspBrId);
            meh->setHitID(icdch);
            if(fConstantErrors) 
                meh->setErrorXYZ(fEx, fEy, fEz);
        }
        else {
            fRecoHitMap[ih] = new CdcSpacepoint(ih, _weightedPlaneConstruction, _cutCov);
            CdcSpacepoint* meh = (CdcSpacepoint*) fRecoHitMap[ih];
            meh->setHitBranchID(fPspBrId);
            meh->setHitID(icdch);
            if(fConstantErrors) 
                meh->setErrorXYZ(fEx, fEy, fEz);
        }
    }
    
    
    // Get the CDC tracks 
    unsigned int nCdcTr = fTrackArray->GetEntriesFast();   
    
    //loop over CdcTracks
    for(unsigned int icdctrk=0; icdctrk<nCdcTr; icdctrk++){
        CdcTrack* itr = (CdcTrack*) (*fTrackArray)[icdctrk];
        double cdcMom = std::fabs(itr->GetMom());
        double cdcTheta = itr->GetTheta();
        double cdcz0 = itr->GetZ0();
        double mx = itr->GetMx();
        double my = itr->GetMy();
        TVector3 cent(mx,my,0.);
        
        // Get the ID of hits and sort by radius
        std::vector<unsigned int> hitIDs = itr->GetCdcIdHits();  
        std::vector<CdcHit*> hitList;
        hitList.reserve(hitIDs.size());
        
        
        for(unsigned int iHit=0; iHit<hitIDs.size(); iHit++){
            CdcHit* ih = (CdcHit*)(*fHitArray)[(hitIDs[iHit])];
            hitList.push_back(ih);
        }
        //sort the hits for GFCand creation
        std::sort(hitList.begin(),hitList.end(),sortingbyrad);
        
        //construct mom-vec for track-rep init: from first two hits
        TVector3 start = (hitList[0])->GetHitPos();
        
        //I choose the stop point at a fifth of the track to get some lever arm
        unsigned int index = 0.2*hitList.size() + 1;
        while(index > hitList.size()-1)
            index--;
        TVector3 stop = (hitList[index])->GetHitPos();
        
        TVector3 radVec = cent - start;
        radVec.RotateZ(0.5*TMath::Pi());
        TVector3 check = stop-start ;   //used for direction setting
        TVector3 check2=check;
        check2.SetZ(0);
        TVector3 radVec2=radVec;
        radVec2.SetZ(0);
        if(((check*radVec<0) and (check2*radVec2>0)) or
              ((check*radVec>0) and (check2*radVec2<0))  ){
            
            std::cout<<" ######################################################## "
                    <<"check*radvec:"<<check*radVec<<", check2d*radvec: "<<check2*radVec<<std::endl;
        }
        if(check2*radVec2 < 0){
            radVec.RotateZ(TMath::Pi());
        }
        TVector3 mom = radVec;
        mom.SetZ(0.);
        mom.SetMag(cdcMom);
        mom.SetTheta(cdcTheta);
        
        TVector3 momerr = mom*0.15;   //15 % error seems ok
        TVector3 poserr(0.03,0.03,6.0); //CDC "pseudo-SP" errors
        
        //Particle ID
        Float_t cdcmass=itr->GetMass();
        Float_t cdccharge=itr->GetCharge();
        int pdg = 211;
        if(cdccharge<0){
            pdg=-211;
        }
        //
        // if requested, the CDC mass calculation is used for assigning pdg.
        // TO BE RECHECKED
        if(fPDG){
            if(cdccharge*cdcmass<-fPiKaonThresh){
                pdg = -321; //k-
            }
            else if(cdccharge*cdcmass<0){
                pdg = -211; // pi-
            }
            else if(!(cdccharge*cdcmass>fPiKaonThresh)){
                pdg = 211; // pi+
            }
            else if(cdccharge*cdcmass<fKaonProtThresh){
                pdg = 321; // k+
            }
            else {
                pdg = 2212; // p
            }
            //
        }
        
        //Set up track representation
        RKTrackRep* rkrep = new RKTrackRep(start, mom, poserr, momerr, pdg);
        //needed to set the time to zero at IP/z-axis
        TVector3 poca,dirInPoca,poca_onwire;
	if(fExtrapToZ){
	  try{
            rkrep->extrapolateToLine(TVector3(0,0,-180),TVector3(0,0,180),poca,dirInPoca,poca_onwire);
            GFDetPlane dp(poca,dirInPoca);
            ((GFAbsTrackRep*)rkrep)->extrapolate(dp);
	  }catch(GFException &exp){
            std::cerr<<"FopiTrackInitTask::exec() couldnt extrapolate rep to z-axis, time set to zero at first point;"<<std::endl;
            std::cerr<<exp.what()<<std::endl;
	  }
	}
        rkrep->setTime(0);
        //create and store GFTrack in output array
        unsigned int ngf = fOutArray->GetEntriesFast();
        //assure that indexing of GFtrack array EQUALS the CdcTrack array 
        assert(ngf == icdctrk);
        GFTrack* gftrk = new((*fOutArray)[ngf]) GFTrack(rkrep);
        GFTrackCand cand;
        //Smoothing
        if(fSmoothing) 
            gftrk->setSmoothing(true);
        else
            gftrk->setSmoothing(false);
        
        for(unsigned int iHit=0; iHit<hitList.size(); iHit++){
            GFAbsRecoHit* thehit = fRecoHitMap[(hitList[iHit])];
            fToWrite.insert(thehit);
            unsigned int nOut = fRHOutArray->GetEntriesFast();
            cand.addHit(fPspBrId,nOut);
            if(fUsePseudoSP) {
                PseudoSpacePoint* cp = (PseudoSpacePoint*) thehit;
                new ((*fRHOutArray)[fRHOutArray->GetEntriesFast()]) PseudoSpacePoint(*cp);
            }
            else {
                CdcSpacepoint* cp = (CdcSpacepoint*) thehit;
                new ((*fRHOutArray)[fRHOutArray->GetEntriesFast()]) CdcSpacepoint(*cp);
            }
        }
        
        
        gftrk->setCandidate(cand);  //Track takes ownership
    }// end loop over CDCTracks
    
    //write out all other recoHits if required, clean up:
    std::map<CdcHit*, GFAbsRecoHit*>::const_iterator it;
    for(it=fRecoHitMap.begin(); it!=fRecoHitMap.end(); it++) {
        GFAbsRecoHit* iHit = it->second;
        if(!fToWrite.count(iHit)) {
            if(fArhPersistence) {
                if(fUsePseudoSP) {
                    PseudoSpacePoint* cp = (PseudoSpacePoint*) iHit;
                    new ((*fRHOutArray)[fRHOutArray->GetEntriesFast()]) PseudoSpacePoint(*cp);
                }
                else {
                    CdcSpacepoint* cp = (CdcSpacepoint*) iHit;
                    new ((*fRHOutArray)[fRHOutArray->GetEntriesFast()]) CdcSpacepoint(*cp);
                }
            }
        }
        delete iHit;  //all of them have been copied during TCA write ...
    }
    
    
    
    if(fVerbose) std::cout<<"    ... "<<nCdcTr<<" CDC tracks setup.\n";
}

void
FopiTrackInitTask::SetConstantPointErrors(double x, double y, double z) {
    fConstantErrors = kTRUE;
    fEx = x;
    fEy = y;
    fEz = z;  
}
ClassImp(FopiTrackInitTask)