//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      RPC and Barrel hits are added to existing TPC-CDC 
//      GFTracks
//      Sept. 2014
//      
//
// Environment:
//      Software developed for the Prototype Detector at FOPI
//
// Author List:
//      Felix Boehmer     (TUM)
// 
//
//-----------------------------------------------------------


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

#include <vector>
#include <cmath>
#include <algorithm>

#include "TpcDigiPar.h"
#include "MatchingTuple.h"
#include "CdcTrack.h"
#include "RpcPixHit.h"
#include "BarPixHit.h"

#include "TClonesArray.h"
#include "TMatrixDSym.h"
#include "GFTrack.h"
#include "RecoHits/GFAbsRecoHit.h"
#include "GFException.h"
#include "RKTrackRep.h"

#include "FairRootManager.h"
#include "FairRuntimeDb.h"
#include "FairRun.h"

ClassImp(FopiRpcBarMergingTask)

FopiRpcBarMergingTask::FopiRpcBarMergingTask()
:  fPersistence(kTRUE),
  fOutBranchName("FopiTrackPreFit"),
  fCdcOutBranchName("CdcTrackPreFit"),
  fTrackBranchName("FopiTrackPostTempFit"),
  fCdcTrackBranchName("CdcTrack"),
  fCdcGFTrackBranchName("CdcTrackPreTempFit"),
  fTupleBranchName("FopiTrackTuples"),
  fRpcHitBranchName("RpcPixHit"),
  fBarHitBranchName("BarPixHit"),
  fTupleArray(NULL),
  fTrackArray(NULL),	 
  fCdcTrackArray(NULL),	 
  fCdcGFTrackArray(NULL), 
  fRpcHitArray(NULL),     
  fBarHitArray(NULL),     
  fTrackOutArray(NULL),     
  fCdcTrackOutArray(NULL),
  fSmoothing(false),
  fDisRpc(kFALSE),
  fPar(NULL),
  fRpcResMean(0),
  fRpcResSig(0.),
  fSigMult(2.),
  fExtrapToZ(kFALSE)
{
    ;
}




FopiRpcBarMergingTask::~FopiRpcBarMergingTask() {
    ;
}




InitStatus
FopiRpcBarMergingTask::Init() {
    
    //Get ROOT Manager ------------------
    FairRootManager* ioman= FairRootManager::Instance();
    if(ioman==0)
    {
        Error("FopiRpcBarMergingTask::Init","RootManager not instantiated!");
        return kERROR;
    }
    
    // Get input ------------------------
    fCdcTrackArray=(TClonesArray*) ioman->GetObject(fCdcTrackBranchName);
    if(fCdcTrackArray==0) {
        Error("FopiRpcBarMergingTask::Init","CDC track array not found!");
        return kERROR;
    }  
    fCdcGFTrackArray=(TClonesArray*) ioman->GetObject(fCdcGFTrackBranchName);
    if(fCdcGFTrackArray==0) {
        Error("FopiRpcBarMergingTask::Init","CDC track array not found!");
        return kERROR;
    }  
    fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);
    if(fTrackArray==0) {
        Error("FopiRpcBarMergingTask::Init","TPC-CDC GFTrack array not found!");
        return kERROR;
    }  
    fTupleArray=(TClonesArray*) ioman->GetObject(fTupleBranchName);
    if(fTupleArray==0) {
        Error("FopiRpcBarMergingTask::Init","TPC-CDC Tuple array not found!");
        return kERROR;
    }  
    fRpcHitArray=(TClonesArray*) ioman->GetObject(fRpcHitBranchName);
    if(fRpcHitArray==0) {
        Error("FopiRpcBarMergingTask::Init","RpcHit array not found!");
        return kERROR;
    }  
    fBarHitArray=(TClonesArray*) ioman->GetObject(fBarHitBranchName);
    if(fBarHitArray==0) {
        Error("FopiRpcBarMergingTask::Init","BarHit array not found!");
        return kERROR;
    }  
    
    
    fTrackOutArray = new TClonesArray("GFTrack");
    fCdcTrackOutArray = new TClonesArray("GFTrack");
    ioman->Register(fOutBranchName,"Tpc",fTrackOutArray,fPersistence);
    ioman->Register(fCdcOutBranchName,"Tpc",fCdcTrackOutArray,fPersistence);
    
    return kSUCCESS;
}

void 
FopiRpcBarMergingTask::SetParContainers() {
    
    //  std::cout<<"TpcResidualTask::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
FopiRpcBarMergingTask::Exec(Option_t* opt) {
    
    std::cout<<"FopiRpcBarMergingTask::Exec()"<<std::endl;
    
    if(fTrackOutArray==NULL)
        Fatal("FopiRpcBarMergingTask::Exec()","No FOPI output array");
    fTrackOutArray->Delete();
    
    if(fCdcTrackOutArray==NULL)
        Fatal("FopiRpcBarMergingTask::Exec()","No CDC output array");
    fCdcTrackOutArray->Delete();
    
    
    std::map<unsigned int, MatchingTuple*> tupleMap;
    
    unsigned int nTuples = fTupleArray->GetEntriesFast();
    for(unsigned int ituple=0; ituple<nTuples; ituple++)  {
        MatchingTuple* it = (MatchingTuple*)fTupleArray->At(ituple);
        //ID of the merged track
        int refID = it->getRefID();
        if(tupleMap.count(refID))
            Fatal("FopiRpcBarMergingTask::Exec","refIDs are not unique!!!");
        tupleMap[refID] = it;
    }
    
    
    //for(unsigned int r=0; r<refIDs.size(); r++) {
    for(unsigned int nt=0; nt<fTrackArray->GetEntriesFast(); nt++) {
        
        GFTrack* ifopi = (GFTrack*)fTrackArray->At(nt);
        
        //make output track
        TVector3 pos, mom;
        TMatrixDSym dummy;
        int pdg = ((RKTrackRep*)(ifopi->getCardinalRep()))->getPDG();
        ifopi->getPosMomCov(pos,mom,dummy);
        TVector3 momerr = mom*0.15;     //15 % error seems ok
        TVector3 poserr(0.03,0.03,6.0); //CDC "pseudo-SP" errors
        
        RKTrackRep* rkrep = new RKTrackRep(pos, mom, poserr, momerr, pdg);
        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<<"FopiRpcBarMergingTask::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);
        GFTrack* outTrack = new((*fTrackOutArray)[nt]) GFTrack(rkrep);
        
        outTrack->setSmoothing(fSmoothing);
        
        GFTrackCand cand(ifopi->getCand());
        
        std::map<unsigned int, MatchingTuple*>::iterator iter;
        //check if we have a corresponding tuple
        iter = tupleMap.find(nt);
        if(iter==tupleMap.end()) {
            outTrack->setCandidate(cand);
            continue;  //just write out
        }
        
        MatchingTuple* it = iter->second;
        //ID of the merged track
        int refID = it->getRefID();
        //consistency check
        assert(refID==nt);
        
        
        //unmatched tracks are of no concern, just copy
        if(!it->isMatched()) {
            outTrack->setCandidate(cand);
            continue;   
        }
        
        //check for matched Barrel/RPC
        int barID, rpcID, cdcID;
        
        //ID of the associated CdcTrack
        try {
            cdcID = it->getIdFromBranch(fCdcTrackBranchName);
        }
        catch(const std::exception& ex) {
            std::cout<<ex.what();
            Fatal("FopiRpcBarMergingTask::Exec","Could not get cdcID from Tuple!");
            continue;
        }
        
        CdcTrack* icdc = (CdcTrack*) fCdcTrackArray->At(cdcID);
        barID = icdc->GetBarHitID();
        rpcID = icdc->GetRpcHitID();
        
        
        
        if(barID < 0 && rpcID < 0) {  //no match, just write out
            outTrack->setCandidate(cand);
            continue;
        }
        
        if(barID >= 0 && rpcID >= 0) { //double match, just write out
            outTrack->setCandidate(cand);
            continue;
        }
        
        
        if(barID >=0) {
            BarPixHit* ibar = (BarPixHit*) fBarHitArray->At(barID);
            if(ibar==NULL) 
                Fatal("FopiRpcBarMergingTask::Exec","Could not get valid Barrel hit!");
            int brId = ibar->getHitBranchID();
            int hitId = ibar->getHitID();
            cand.addHit(brId,hitId);
        }
        
        if(rpcID >=0) {
            RpcPixHit* irpc = (RpcPixHit*) fRpcHitArray->At(rpcID);
            if(irpc==NULL) 
                Fatal("FopiRpcBarMergingTask::Exec","Could not get valid RPC hit!");
            int brId = irpc->getHitBranchID();
            int hitId = irpc->getHitID();
            
            //ADDITIONAL Z-RESIDUAL CHECK:
            if(fDisRpc) {
                double zRes;
                try {
                    GFDetPlane lastPlane = irpc->getDetPlane(ifopi->getCardinalRep());
                    TVector3 exPos = ifopi->getCardinalRep()->getPos(lastPlane);
                    TVector3 residual = (lastPlane.getO() - exPos);
                    zRes = residual.Z();
                }
                catch(const GFException& ex) {
                    std::cout<<"FopiRpcBarMergingTask::Exec() RPC extrapolation failed; not adding the hit.";
                    std::cout<<ex.what();
                    std::cout<<std::endl;
                    outTrack->setCandidate(cand);
                    continue;
                }
                if(std::fabs(zRes-fRpcResMean) < fSigMult*fRpcResSig) {
                    cand.addHit(brId,hitId);
                }
                else  {
                    std::cout<<"FopiRpcBarMergingTask::Exec() RPC extrapolation outside valid window. Not adding RPC hit";
                    std::cout<<std::endl;
                    outTrack->setCandidate(cand);
                    continue;
                }
            }
            else {
                cand.addHit(brId,hitId);
            }
        }
        
        //we only get here if everything checked out, otherwise a copy of the original track will be written
        outTrack->setCandidate(cand);
        if(barID >=0) 
            std::cout<<"    FopiRpcBarMergingTask::Exec() Bar hit added"<<std::endl;
        if(rpcID >=0) 
            std::cout<<"    FopiRpcBarMergingTask::Exec() Rpc hit added"<<std::endl;
        
        
    }
    
    //now match CDC-only tracks. no check on residual is performed
    for(unsigned int ncdc=0; ncdc<fCdcTrackArray->GetEntriesFast(); ncdc++) {
        CdcTrack* icdc = (CdcTrack*) fCdcTrackArray->At(ncdc);
        GFTrack* icdcGF = (GFTrack*) fCdcGFTrackArray->At(ncdc);
        GFTrackCand icdcCand(icdcGF->getCand());
        
        //make output track
        TVector3 pos, mom;
        TMatrixDSym dummy;
        int pdg = ((RKTrackRep*)(icdcGF->getCardinalRep()))->getPDG();
        icdcGF->getPosMomCov(pos,mom,dummy);
        TVector3 momerr = mom*0.15;     //15 % error seems ok
        TVector3 poserr(0.03,0.03,6.0); //CDC "pseudo-SP" errors
        
        RKTrackRep* rkrep = new RKTrackRep(pos, mom, poserr, momerr, pdg);
        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<<"FopiRpcBarMergingTask::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);
	GFTrack* icdcGFout = new((*fCdcTrackOutArray)[ncdc]) GFTrack(rkrep);
        
        icdcGFout->setSmoothing(fSmoothing);
        
        
        int barID = icdc->GetBarHitID();
        int rpcID = icdc->GetRpcHitID();
        
        if(barID < 0 && rpcID < 0) {  //no match, just write out
            icdcGFout->setCandidate(icdcCand);
            continue;
        }
        
        if(barID >= 0 && rpcID >= 0) { //double match, just write out
            icdcGFout->setCandidate(icdcCand);
            continue;
        }
        
        if(barID >= 0) {
            BarPixHit* ibar = (BarPixHit*) fBarHitArray->At(barID);
            if(ibar==NULL) 
                Fatal("FopiRpcBarMergingTask::Exec","Could not get valid Barrel hit!");
            int brId = ibar->getHitBranchID();
            int hitId = ibar->getHitID();
            icdcCand.addHit(brId,hitId);
        }
        
        if(rpcID >= 0) {
            RpcPixHit* irpc = (RpcPixHit*) fRpcHitArray->At(rpcID);
            if(irpc==NULL) 
                Fatal("FopiRpcBarMergingTask::Exec","Could not get valid Rpcrel hit!");
            int brId = irpc->getHitBranchID();
            int hitId = irpc->getHitID();
            icdcCand.addHit(brId,hitId);
        }
        
        icdcGFout->setCandidate(icdcCand);
        
    }
    
}