//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//     Analysis task for lambdas using GFRaveVertex
//     to build Lambdas. No PID assumed
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sverre Doerheim
//
//-----------------------------------------------------------
#include "FopiLambdaAnaTask3.h"

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "FopiLambdaCand.h"
#include "TpcVertex.h"
#include "MatchingPair.h"
#include "CdcTrack.h"

// GENFIT headers
#include "GFTrack.h"
#include "GFRaveVertex.h"
#include "GFRaveTrackParameters.h"
#include "GFException.h"

//ROOT headers
#include "TDatabasePDG.h"
#include "TClonesArray.h"
#include "TLorentzVector.h"

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


FopiLambdaAnaTask3::FopiLambdaAnaTask3()
  :FairTask("FopiLambda3"),
  fPersistence(kTRUE),
  fTrackBranchName("TpcCdcTrack"),
  fVtxBranchName("LambdaVertex"),
  fRaveVtxBranchName("LambdaRaveVertex"),
  fCandName("LambdaCandRave"),
  fFopiTupleBranchName("FopiTrackTuples"),
  fCdcGFTrackBranchName("CdcTrackPostFit"),
  fCdcTrackBranchName("CdcTrack"),
  fMcPid(false)
{

}

InitStatus FopiLambdaAnaTask3::Init(){
  if(fVerbose>0){
    std::cerr<<"FopiLambdaAnaTask3::Init()"<<"\n";
  }
  FairRootManager* ioman= FairRootManager::Instance();
  if(ioman==0) {
    Error("FopiLambdaAnaTask2::Init","RootManager not instantiated!");
    return kERROR;
  }
  fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);
  if(fTrackArray==0) {
    Error("FopiLambdaAnaTask3::Init","track-array not found!");
    return kERROR;
    assert(0);
  }
  fVtxArray=(TClonesArray*) ioman->GetObject(fVtxBranchName);
  if(fVtxArray==0) {
    Error("FopiLambdaAnaTask3::Init","vtx-array not found!");
    return kERROR;
    assert(0);
  }
  fRaveVtxArray=(TClonesArray*) ioman->GetObject(fRaveVtxBranchName);
  if(fRaveVtxArray==0) {
    Error("FopiLambdaAnaTask3::Init","vtx-array not found!");
    return kERROR;
    assert(0);
  }
  std::cerr<<"getting fFopiTupleArray with name: "<<fFopiTupleBranchName<<"\n";
  fFopiTupleArray=(TClonesArray*) ioman->GetObject(fFopiTupleBranchName);
  if(fFopiTupleArray==0)    {
    std::cerr<<"FopiLambdaAnaTask3::Init() no matching array found, can not use PID"<<"\n";
  }
  std::cerr<<"getting CdcTrack Branch with name: "<<fCdcTrackBranchName<<"\n";
  fCdcTrackArray=(TClonesArray*) ioman->GetObject(fCdcTrackBranchName);
  if(fCdcTrackArray==0)    {
    std::cerr<<"FopiLambdaAnaTask3::Init() no CdcTrack array found, can not use PID"<<"\n";
  }
  std::cerr<<"creating tclonesarray out"<<"\n";
  fLambdaOutArray = new TClonesArray("FopiLambdaCand");
  std::cerr<<"registering outbranch"<<"\n";
  ioman->Register(fCandName,"Tpc",fLambdaOutArray,fPersistence);
  fPDG=TDatabasePDG::Instance();
  std::cerr<<"FopiLambdaAnaTask3::init() done"<<"\n";
  return kSUCCESS;

}

void FopiLambdaAnaTask3::Exec(Option_t* opt) {
  //std::cerr<<"FopiLambdaAnaTask3::Exec()"<<fVerbose<<" #############################################################################################"<<"\n";
  fLambdaOutArray->Delete();
  unsigned int nVtx=fVtxArray->GetEntriesFast();
  unsigned int nVtx2=fRaveVtxArray->GetEntriesFast();
  //loop over vertices, Lambda candidates
  //std::cout<<"nVtx(tpc,rave) "<<nVtx<<", "<<nVtx2<<"\n";
  for(unsigned int iVtx=0;iVtx<nVtx;++iVtx){
    TpcVertex* theVtx=(TpcVertex*)(*fVtxArray)[iVtx];
    GFRaveVertex* theVtx2=(GFRaveVertex*)(*fRaveVtxArray)[iVtx];
//    std::cout<<"nTrk: "<<theVtx->getNTracks()<<", "
//        <<theVtx2->getNTracks()<<"\n";
    if(theVtx->getNTracks()!=2){
      std::cout<<"ntracks!=2, "<<theVtx->getNTracks()<<"\n";
      continue;
    }
    //if PID available, extract from CdcTrack
    double track1Mass=0.;
    double track2Mass=0.;
    if(fCdcTrackArray!=0&&fFopiTupleArray!=0){
      int trackId1=theVtx->getTracks().at(0);
      MatchingTuple* mp = (MatchingTuple*) fFopiTupleArray->At(trackId1);
      //is CDCinformation available?
      if(true){
        //std::cout<<"got matching tuple"<<"\n";
      }
      if(mp->isMatched()){
        if(true){
          //std::cout<<"isMatched"<<"\n";
//          for(unsigned int test=0;test<mp->getListOfBranchNames().size();++test){
//            std::cout<<mp->getListOfBranchNames()[test]<<"\n";
//          }
        }
        if(mp->hasBranch(fCdcTrackBranchName)){
          std::cout<<"found branch"<<"\n";
          int cdcTrackId=mp->getIdFromBranch(fCdcTrackBranchName);
          CdcTrack* tr=(CdcTrack*)fCdcTrackArray->At(cdcTrackId);
          track1Mass=tr->GetMass();
       //   std::cout<<"track1Mass "<<track1Mass<<"\n";
        }
      }
      int trackId2=theVtx->getTracks().at(1);
      mp = (MatchingTuple*) fFopiTupleArray->At(trackId2);
      //is CDCinformation available?
      if(mp->isMatched()){
        if(mp->hasBranch(fCdcTrackBranchName)){
          int cdcTrackId=mp->getIdFromBranch(fCdcTrackBranchName);
          CdcTrack* tr=(CdcTrack*)fCdcTrackArray->At(cdcTrackId);
          track2Mass=tr->GetMass();
         // std::cout<<"track2Mass "<<track2Mass<<"\n";
        }
      }
    }

    GFRaveTrackParameters* par1;
    GFRaveTrackParameters* par2;
    try{
      par1 = theVtx2->getParameters(0);
      par2 = theVtx2->getParameters(1);
    }catch(GFException &exp){
      std::cerr<<exp.what()<<"\n";
      continue;
    }
    //  std::cout<<"Vtx nr: "<<iVtx<<"\n";
    GFTrack* track1=par1->getTrack();
    GFTrack* track2=par2->getTrack();

    // std::cerr<<"TrackMom1: "<<track1->getMom().Mag()
    //     <<" VtxTrackParMom1: "<<par1->getMom().Mag()<<"\n";
    // std::cerr<<"TrackMom2: "<<track2->getMom().Mag()
    //    <<" VtxTrackParMom2: "<<par2->getMom().Mag()<<"\n";

    try{
      if(fMcPid){
        if(abs(track1->getCardinalRep()->getPDG())==2212 &&
           abs(track2->getCardinalRep()->getPDG())==211){
          TLorentzVector proton;
          proton.SetVectM(par1->getMom(),fPDG->GetParticle(2212)->Mass());
          TLorentzVector pion;
          pion.SetVectM(par2->getMom(),fPDG->GetParticle(211)->Mass());
          TLorentzVector lambda=pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          cand->setPMom(par1->getMom());
          cand->setPiMom(par2->getMom());
          cand->setPCdcMass(track1Mass);
          cand->setPiCdcMass(track2Mass);

        }else if(abs(track2->getCardinalRep()->getPDG())==2212 &&
                 abs(track1->getCardinalRep()->getPDG())==211){
          TLorentzVector proton;
          proton.SetVectM(par2->getMom(),fPDG->GetParticle(2212)->Mass());
          TLorentzVector pion;
          pion.SetVectM(par1->getMom(),fPDG->GetParticle(211)->Mass());
          TLorentzVector lambda=pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          cand->setPMom(par2->getMom());
          cand->setPiMom(par1->getMom());
          cand->setPCdcMass(track2Mass);
          cand->setPiCdcMass(track1Mass);
          if(par1->hasTrack()){}
          // cand->setPitrack(par1->getTrack());
          if(par2->hasTrack()){}
          //  cand->setPtrack(par2->getTrack());

        }

      }else{
        if(par1->getCharge() > 0 &&
           par2->getCharge() < 0){
          std::cout<<"vertex # "<<iVtx<<", charge 1>0, charge2<0 \n";
          TLorentzVector proton;
          proton.SetVectM(par1->getMom(),fPDG->GetParticle(2212)->Mass());
          TLorentzVector pion;
          pion.SetVectM(par2->getMom(),fPDG->GetParticle(211)->Mass());
          TLorentzVector lambda=pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          cand->setPMom(par1->getMom());
          cand->setPiMom(par2->getMom());
          cand->setPCdcMass(track1Mass);
          cand->setPiCdcMass(track2Mass);


        }else if(par1->getCharge() < 0 &&
                 par2->getCharge() > 0){
          std::cout<<"anatask3 test2 "<<iVtx<<", charge 1>0, charge2<0 \n";
          TLorentzVector proton;
          proton.SetVectM(par2->getMom(),fPDG->GetParticle(2212)->Mass());
          TLorentzVector pion;
          pion.SetVectM(par1->getMom(),fPDG->GetParticle(211)->Mass());
          TLorentzVector lambda=pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          cand->setPMom(par2->getMom());
          cand->setPiMom(par1->getMom());
          cand->setPCdcMass(track2Mass);
          cand->setPiCdcMass(track1Mass);
          if(par1->hasTrack()){}
          // cand->setPitrack(par1->getTrack());
          if(par2->hasTrack()){}
          //  cand->setPtrack(par2->getTrack());

        }
      }
    }catch(GFException &exp){
      std::cerr<<exp.what()<<"\n";
    }
  }
}