//-----------------------------------------------------------
// 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 "FopiLambdaAnaTask2.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;


FopiLambdaAnaTask2::FopiLambdaAnaTask2()
  : FairTask("FopiLambda"), 
    fPersistence(kTRUE),
    fVtxUseVacuumPropagator(false),
    fVtxMethod("kalman-smoothing:1")
{
  fTrackBranchName = "TrackPostFitComplete";
  fVtxOutBranchName = "GFVertex";
  fCandName = "FopiLambdaCand2";
  fVerbose = 0;
} 


FopiLambdaAnaTask2::~FopiLambdaAnaTask2()
{

}

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

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


void
FopiLambdaAnaTask2::Exec(Option_t* opt) {
  fLambdaOutArray->Delete();
  fVtxOutArray->Delete();
  // if(nEv%1000==0&&fVerbose>0)
  std::cerr<<"FopiLambdaAnaTask2::Exec "<<nEv<<"########################################"<<std::endl;
  nEv++;
  unsigned int ntracks = fTrackArray->GetEntriesFast();
  vector<GFTrack*> protons;
  vector<GFTrack*> pi_min;


  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();
  

  for(unsigned int iTr=0;iTr<ntracks;++iTr){
    GFTrack* track = (GFTrack*) (*fTrackArray)[iTr];
    GFAbsTrackRep* rep = track->getCardinalRep();
    try{
      switch(rep->getPDG()){
      case 2212:
        protons.push_back(track);
        break;
      case -211:
        pi_min.push_back(track);
        break;
      default:
        break;
      }
    }catch(GFException &exp){
      cout<<exp.what()<<endl;
    }
  }

  std::cerr<<"track arrays filled"<<endl;
  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;
    }
  }

  std::cerr<<"track buffers for vtx filled"<<endl;
  for(int iVtx=0;iVtx<vtxCounter;++iVtx){
    // if(iVtx>1)
    //   break;
    for(int i=0;i<fVertexBuffer->size();++i){
      delete fVertexBuffer->at(i);
    }
    fVertexBuffer->clear();
    cerr<<"clearing vertex buffer"<<endl;
    try{
      fVertexFactory->findVertices(fVertexBuffer, fVtxTrackBuffer.at(iVtx), false);
    }catch(GFException &exp){
      cout<<exp.what()<<endl;
      continue;
    }
    GFRaveVertex* theVtx;
    if(fVertexBuffer->size()>0){
      std::cerr<<"vtxBuffSize"<<fVertexBuffer->size()<<std::endl;
      theVtx=fVertexBuffer->back();//at(iVtx);
    }else{
      continue;
    }
    
    double ndf=theVtx->getNdf();
    double chi2=theVtx->getChi2();

    int nTrk=theVtx->getNTracks();
    TVector3 dummy(0,0,0);
    if(nTrk==2){
      GFRaveTrackParameters* par1;
      GFRaveTrackParameters* par2;
      try{
        par1 = theVtx->getParameters(0);
        par2 = theVtx->getParameters(1);
      }catch(GFException &exp){
        cout<<exp.what()<<endl;
        continue;
      }
      try{
        if((par1->getPdg()==-211&&par2->getPdg()==2212)){
          
          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());
          
          GFRaveVertex* outVtx=new((*fVtxOutArray)[fVtxOutArray->GetEntriesFast()])GFRaveVertex(*theVtx);
          //          outVtx->Print();
          //  (*fVtxOutArray)[fVtxOutArray->GetLast()]->Print();
          //          cand->setRaveVertex(outVtx);
          std::cerr<<"par1=-211********************WRITING LAMBDA CAND AND VERTEX nr:"<<iVtx<<"********************"<<std::endl;
                    
        } else if (par2->getPdg()==-211&&par1->getPdg()==2212){
          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());
          GFRaveVertex* outVtx=new((*fVtxOutArray)[fVtxOutArray->GetEntriesFast()])GFRaveVertex(*theVtx);
          //outVtx->Print();
       
          //          cand->setRaveVertex(outVtx);
          
          std::cerr<<"par2=-211********************WRITING LAMBDA CAND AND VERTEX nr:"<<iVtx<<"********************"<<std::endl;
        }
      }catch(GFException &exp){
        cout<<exp.what()<<endl;
        continue;
      }
    }
  }
  // for(unsigned int i=0;i<fVtxOutArray->GetEntriesFast();++i){
  //   (*fVtxOutArray)[i]->Print();
  // }
  cerr<<"FopiLambdaAnaTask2::exec end"<<endl;
  return;
}