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

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

// This Class' Header ------------------
#include "FopiLambdaAnaTask_Paul.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 "TH1D.h"
#include "TH2D.h"
#include "TH2I.h"
#include "TLorentzVector.h"
#include "TDatabasePDG.h"

// Class Member definitions -----------
using namespace std;


FopiLambdaAnaTask_Paul::FopiLambdaAnaTask_Paul()
  : fPersistence(kTRUE),
    fVtxUseVacuumPropagator(false),
    fVtxMethod("kalman-smoothing:1")
{
  fTrackBranchName = "TpcHelPostFit";
  fVtxOutBranchName = "GFVertex";
  fCandName = "TpcHeliLambdaCand";
  fVerbose = 0;
} 


FopiLambdaAnaTask_Paul::~FopiLambdaAnaTask_Paul()
{

}

InitStatus
FopiLambdaAnaTask_Paul::Init(){
  if(fVerbose>0){
    std::cerr<<"FopiLambdaAnaTask_Paul::Init()"<<std::endl;
  }
  
  nEv=0;
  
  //Get ROOT Manager ###########################################################
  FairRootManager* ioman= FairRootManager::Instance();
  if(ioman==0) {
    Error("FopiLambdaAnaTask_Paul::Init","RootManager not instantiated!");
    return kERROR;
  }
  
  FairField* field=FairRunAna::Instance()->GetField();
  GFFieldManager::getInstance()->init(new PndFieldAdaptor(field));
 
  fTrackArray=(TClonesArray*) ioman->GetObject(fTrackBranchName);
  if(fTrackArray==0) {
    Error("FopiLambdaAnaTask_Paul::Init","track-array not found!");
    return kERROR;
    assert(0);
  }
  
  // register new branches
  fLambdaOutArray = new TClonesArray("FopiLambdaCand");
  fVtxOutArray = new TClonesArray("GFRaveVertex");
  ioman->Register(fCandName,"Tpc",fLambdaOutArray,fPersistence);
  ioman->Register(fVtxOutBranchName,"Tpc",fVtxOutArray,fPersistence);

  // prepare buffers
  fVertexFactory = new GFRaveVertexFactory(fVerbose, fVtxUseVacuumPropagator);
  fVertexFactory->setMethod(fVtxMethod);
  fVertexBuffer = new std::vector < GFRaveVertex* >;
  fPDG=TDatabasePDG::Instance();
  
  return kSUCCESS;
}


void
FopiLambdaAnaTask_Paul::Exec(Option_t* opt) {
  if (fVerbose>0)
    std::cerr<<"FopiLambdaAnaTask_Paul::Exec "<<nEv<<"########################################"<<std::endl;
  nEv++;

  // initialise list of proton and pi- candidates
  vector<GFTrack*> protons;
  vector<GFTrack*> pi_min;
  
  // clean-up output buffers
  fLambdaOutArray->Delete();
  fVtxOutArray->Delete();

  // number of available tracks
  unsigned int ntracks = fTrackArray->GetEntriesFast();

  // clean-up vertex buffers
  for(unsigned int i=0;i<fVtxTrackBuffer.size();++i){
    fVtxTrackBuffer.at(i).clear();
  }
  fVtxTrackBuffer.clear();
  
  // loop over tracks and fill list of proton and pi- candidates
  // all positive particles are assumed to be protons
  // all negative particles are considered to be pi-
  // only trakcs are considered with proper fit status
  for(unsigned int iTr=0;iTr<ntracks;++iTr){
    GFTrack* track = (GFTrack*) (*fTrackArray)[iTr];
    GFAbsTrackRep* rep = track->getCardinalRep();
    if (rep->getStatusFlag() == 0) {
      try{
        if (rep->getCharge() > 0.) {
          protons.push_back(track);
        } else {
          pi_min.push_back(track);
        }
      } catch(GFException &exp){
        cout<<exp.what()<<endl;
      }
    }
  }

  // fill list of proton-pion combinations
  int vtxCounter=0;
  for(unsigned int iP=0;iP<protons.size();++iP){
    for(unsigned int iPi=0;iPi<pi_min.size();++iPi){
      std::vector<GFTrack*> a;
      a.reserve(2);
      a.push_back(protons.at(iP));
      a.push_back(pi_min.at(iPi));
      fVtxTrackBuffer.push_back(a);
      ++vtxCounter;
    }
  }

  // loop over all proton-pion combinations
  for(int iVtx=0;iVtx<vtxCounter;++iVtx){
    
    // make sure that vertex buffer is cleared
    for(int i=0;i<fVertexBuffer->size();++i){
      delete fVertexBuffer->at(i);
    }
    fVertexBuffer->clear();

    // do the actual vertexing
    try {
      fVertexFactory->findVertices(fVertexBuffer, fVtxTrackBuffer.at(iVtx), false);
    } catch(GFException &exp) {
      cout<<exp.what()<<endl;
      continue;
    }
    GFRaveVertex* theVtx;
    if (fVertexBuffer->size()>0) {
      theVtx=fVertexBuffer->back();
    } else {
      continue;
    }
    
    double ndf=theVtx->getNdf();
    double chi2=theVtx->getChi2();

    // construct invariant mass
    int nTrk=theVtx->getNTracks();
    TVector3 dummy(0,0,0);
    if (nTrk == 2) {                // there should always be 2 tracks
      GFRaveTrackParameters* par1;
      GFRaveTrackParameters* par2;
      try {
        par1 = theVtx->getParameters(0);
        par2 = theVtx->getParameters(1);
      } catch(GFException &exp){
        cout<<exp.what()<<endl;
        continue;
      }
      try { 
        TLorentzVector proton;
        TLorentzVector pion;
        TLorentzVector lambda;
        if ((par1->getCharge()<0.&&par2->getCharge()>0.)) {
          
          proton.SetVectM(par2->getMom(),fPDG->GetParticle(2212)->Mass());
          pion.SetVectM(par1->getMom(),fPDG->GetParticle(211)->Mass());
          lambda = pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          
          GFRaveVertex* outVtx=new((*fVtxOutArray)[fVtxOutArray->GetEntriesFast()])GFRaveVertex(*theVtx);
                    
        } else if (par2->getCharge()<0.&&par1->getCharge()>0.){
          
          proton.SetVectM(par1->getMom(),fPDG->GetParticle(2212)->Mass());
          pion.SetVectM(par2->getMom(),fPDG->GetParticle(211)->Mass());
          lambda=pion+proton;
          FopiLambdaCand* cand=new((*fLambdaOutArray)[fLambdaOutArray->GetEntriesFast()])FopiLambdaCand(lambda,theVtx->getPos());
          
          GFRaveVertex* outVtx=new((*fVtxOutArray)[fVtxOutArray->GetEntriesFast()])GFRaveVertex(*theVtx);
          
        }
      } catch (GFException &exp) {
        cout<<exp.what()<<endl;
        continue;
      }
    }
  }

  return;
}