//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//     Analysis task for lambdas
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sverre Doerheim
//
//-----------------------------------------------------------

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

// This Class' Header ------------------
#include "FopiLambdaRaveVtxTask.h"
#include "FopiLambdaCand.h"

// C/C++ Headers ----------------------
#include <iostream>
#include <vector>
#include <map>
#include <assert.h>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "GFDetPlane.h"
#include "GFFieldManager.h"
#include "GFKalman.h"

#include "GFRecoHitFactory.h"
#include "GFException.h"
#include "GFRaveVertexFactory.h"
#include "GFRaveVertex.h"

#include "GFDetPlane.h"

#include "RKTrackRep.h"

#include "PndDetectorList.h"

#include "PndFieldAdaptor.h"
#include "FairField.h"
#include "FairRunAna.h"
#include "GFTrackProximity.h"
#include "GFRaveVertex.h"
#include "TpcVertex.h"
#include "TLorentzVector.h"
#include "TDatabasePDG.h"

// Class Member definitions -----------
using namespace std;
ClassImp(FopiLambdaRaveVtxTask)

FopiLambdaRaveVtxTask::FopiLambdaRaveVtxTask()
    : FairTask("FopiLambdaVtx"),
      fPersistence(kTRUE),
      fVtxUseVacuumPropagator(false),fUseMCID(false),
      fVtxMethod("kalman-smoothing:1")
{
    fTrackBranchName = "TpcCdcTrack";
    fVtxOutBranchName = "LambdaVertex";
    fRaveVtxOutBranchName = "LambdaRaveVertex";
    fVerbose = 0;
} 


FopiLambdaRaveVtxTask::~FopiLambdaRaveVtxTask()
{
    
}

InitStatus
FopiLambdaRaveVtxTask::Init(){
    //if(fVerbose>0){
    std::cerr<<"FopiLambdaRaveVtxTask::Init()"<<"\n";
    //}
    
    //Get ROOT Manager
    FairRootManager* ioman= FairRootManager::Instance();
    if(ioman==0) {
        Error("FopiLambdaRaveVtxTask::Init","RootManager not instantiated!");
        return kERROR;
    }
    
    // Get input collection
    fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);
    if(fTrackArray==0) {
        //    cerr<<fTrackBranchName<<endl;
        Error("FopiLambdaRaveVtxTask::Init","track-array not found!");
        return kERROR;
        assert(0);
    }
    
    
    
    
    nEv=0;
    //Register output arrays
    fVtxOutArray = new TClonesArray("TpcVertex");
    ioman->Register(fVtxOutBranchName,"Tpc",fVtxOutArray,fPersistence);
    fRaveVtxOutArray = new TClonesArray("GFRaveVertex");
    ioman->Register(fRaveVtxOutBranchName,"Tpc",fRaveVtxOutArray,fPersistence);
    
    
    
    return kSUCCESS;
}


void
FopiLambdaRaveVtxTask::Exec(Option_t* opt) {
    //if(nEv%1000==0&&fVerbose>0)
    //std::cerr<<"FopiLambdaRaveVtxTask::Exec "<<nEv<<"########################################"<<"\n"
    //clear output array
    fVtxOutArray->Delete();
    fRaveVtxOutArray->Delete();
    
    fVertexFactory = new GFRaveVertexFactory(fVerbose, fVtxUseVacuumPropagator);
    fVertexFactory->setMethod(fVtxMethod);
    fVertexBuffer = new std::vector < GFRaveVertex* >;
    
    nEv++;
    
    unsigned int ntracks = fTrackArray->GetEntriesFast();
    vector<GFTrack*> tracks_positive;
    vector<GFTrack*> tracks_negative;
    vector<unsigned int> tracks_positive_index;
    vector<unsigned int> tracks_negative_index;
    
    //clear track buffer
    for(unsigned int i=0;i<fVtxTrackBuffer.size();++i){
        fVtxTrackBuffer.at(i).clear();
    }
    fVtxTrackBuffer.clear();
    for(unsigned int i=0;i<fVertexBuffer->size();++i){
        delete fVertexBuffer->at(i);
    }
    fVertexBuffer->clear();
    
    //######################################################################################
    //gather positive and negative tracks
    //std::cout<<"FopiLambdaRAveVtx sorting tracks ######################"<<endl;
    for(unsigned int iTr=0;iTr<ntracks;++iTr){
        GFTrack* track = (GFTrack*) (*fTrackArray)[iTr];
        GFAbsTrackRep* rep = track->getCardinalRep();
        if(rep->getStatusFlag()!=0)
            continue;
        try{
            if(fUseMCID){
                if(abs(rep->getPDG())==2212){
                    tracks_positive.push_back(track);
                    tracks_positive_index.push_back(iTr);
                }
                if(abs(rep->getPDG())==211){
                    tracks_negative.push_back(track);
                    tracks_negative_index.push_back(iTr);
                }
            }else{
                switch((int)(rep->getCharge())){
                case 1:
                    tracks_positive.push_back(track);
                    tracks_positive_index.push_back(iTr);
                    break;
                case -1:
                    tracks_negative.push_back(track);
                    tracks_negative_index.push_back(iTr);
                    break;
                default:
                    break;
                }
            }
            
            fVertexBuffer = new std::vector < GFRaveVertex* >;
            
        }catch(GFException &exp){
            cout<<exp.what()<<endl;
        }
    }
    //if(fVerbose)
    //std::cerr<<"track arrays filled"<<endl;
    //######################################################################################
    
    //######################################################################################
    //build all possible pairs of positive and negative tracks
    std::vector < std::vector< unsigned int > > fVtxTrack_id_Buffer;
    //std::cout<<"FopiLambdaRAveVtx building vtx cand ######################\n";
    //  std::cout<<"pos tracks: "<<tracks_positive.size()<<"\n";
    //  std::cout<<"neg tracks: "<<tracks_negative.size()<<"\n";
    int test=0;
    for(unsigned int iP=0;iP<tracks_positive.size();++iP){
        for(unsigned int iPi=0;iPi<tracks_negative.size();++iPi){
            std::vector<GFTrack*> a;
            std::vector<unsigned int> b;
            GFTrack* track_p=tracks_positive.at(iP);
            unsigned int track_id_p=tracks_positive_index.at(iP);
            
            GFTrack* track_pi=tracks_negative.at(iPi);
            unsigned int track_id_pi=tracks_negative_index.at(iPi);
            
            
            test++;
            a.push_back(track_p);
            a.push_back(track_pi);
            fVtxTrackBuffer.push_back(a);
            
            b.push_back(track_id_p);
            b.push_back(track_id_pi);
            fVtxTrack_id_Buffer.push_back(b);
        }
    }
    if(fVerbose)
        std::cerr<<"track buffers for vtx filled"<<endl;
    //######################################################################################
    
    //######################################################################################
    //loop over track-pairs/vtx-candidates
    int nVtx= fVtxTrackBuffer.size();
    for(int iVtx=0;iVtx<nVtx;++iVtx){
        if(fVerbose){
            cerr<<"clearing vertex buffer"<<endl;
        }
        fVertexBuffer->clear();
        //fit a vertex
        //    std::cout<<"FopiLambdaRaveVtxTask Vtx nr: "<<iVtx<<"\n";
        try{
            fVertexFactory->findVertices(fVertexBuffer, fVtxTrackBuffer.at(iVtx), false);
        }catch(GFException &exp){
            cout<<exp.what()<<endl;
            continue;
        }
        if(fVertexBuffer->size()==0){
            //      std::cout<<"no vtx, tracks:"<<(fVtxTrackBuffer.at(iVtx)).size()<<"\n";
            continue;
        }
        GFRaveVertex* theVtx=fVertexBuffer->back();
        GFRaveTrackParameters* par1;
        GFRaveTrackParameters* par2;
        //std::cerr<<"nVtx "<<fVertexBuffer->size()<<"\n"
        //std::cerr<<"n"<<theVtx->getNTracks()<<" "<<fVtxTrackBuffer.at(iVtx).size()<<"\n"
        try{
            par1 = theVtx->getParameters(0);
            par2 = theVtx->getParameters(1);
        }catch(GFException &exp){
            std::cerr<<exp.what()<<std::endl;
            continue;
        }
        
        
        //std::cerr<<par1->getTrack()<<" "<<fVtxTrackBuffer.at(iVtx).at(0)<<std::endl;
        double charge1=par1->getCharge();
        int pdg1=par1->getPdg();
        //std::cerr<<par2->getTrack()<<" "<<fVtxTrackBuffer.at(iVtx).at(1)<<std::endl;
        double charge2=par2->getCharge();
        int pdg2=par2->getPdg();
        unsigned int id1=0;
        unsigned int id2=0;
        if(fUseMCID){
            if(abs(pdg1)==2212){
                id1=fVtxTrack_id_Buffer.at(iVtx).at(0);
                id2=fVtxTrack_id_Buffer.at(iVtx).at(1);
            }else{
                id1=fVtxTrack_id_Buffer.at(iVtx).at(1);
                id2=fVtxTrack_id_Buffer.at(iVtx).at(0);
            }
        }else{
            if(charge1>0){
                //std::cout<<"charge1 positive, carge2:"<<charge2<<"\n";
                id1=fVtxTrack_id_Buffer.at(iVtx).at(0);
                id2=fVtxTrack_id_Buffer.at(iVtx).at(1);
            }else{
                //std::cout<<"charge1 negative, carge2:"<<charge2<<"\n";
                id1=fVtxTrack_id_Buffer.at(iVtx).at(1);
                id2=fVtxTrack_id_Buffer.at(iVtx).at(0);
            }
        }
        std::vector<unsigned int> ids;
        ids.push_back(id1);
        ids.push_back(id2);
        //debug ############################################################
        
        unsigned int index=fVtxOutArray->GetEntriesFast();
        double dist=(theVtx->getParameters(0)->getPos()-theVtx->getParameters(1)->getPos()).Mag();
        GFRaveVertex* outVtx=new((*fRaveVtxOutArray)[fRaveVtxOutArray->GetEntriesFast()])GFRaveVertex(*theVtx);
        TpcVertex* outVtx2=new((*fVtxOutArray)[fVtxOutArray->GetEntriesFast()])TpcVertex(theVtx->getPos(),
                                                                                         theVtx->getCov(),
                                                                                         dist);
        outVtx2->setRaveIndex(index,"LambdaRaveVertex");
        outVtx2->setIndex(index);
        outVtx2->setTrackIds(ids,fTrackBranchName);
        //}
        
        
    }
    
    // for(unsigned int i=0;i<fVtxOutArray->GetEntriesFast();++i){
    //   (*fVtxOutArray)[i]->Print();
    // }
    delete fVertexFactory;
    if(fVerbose)
        cerr<<"FopiLambdaRaveVtxTask::exec end"<<endl;
    return;
}