#include "evt/Event.hh"
#include "Bbos/bank_interface.hh"

#include "VertexObjects/ZVertexColl.hh"
#include "VertexObjects/VertexColl.hh"
#include "VertexObjects/Vertex.hh"
#include "Stntuple/obj/TStnVertexBlock.hh"
#include "Stntuple/mod/InitStntupleDataBlocks.hh"
#include <Stntuple/mod/StntupleUtilities.hh>

//_____________________________________________________________________________ 
static Int_t StntupleInitVertex(TStnVertex* vtx, const Vertex* hartmut) {
  // initialize TStnVertex starting from Hartmut's vertex

  if (hartmut->type() == Vertex::Primary) {
    vtx->fVClass = TStnVertex::kPRIMARY;
  }
  else if(hartmut->type() == Vertex::Secondary) {
    vtx->fVClass = -2;
  }
  else {
    vtx->fVClass = TStnVertex::kUNCLASSIFIED;
  }

  vtx->fVertex.SetXYZ(hartmut->x(),hartmut->y(),hartmut->z());

  vtx->fMomentum.SetXYZM(hartmut->px(),hartmut->py(),hartmut->pz(),hartmut->mass());

  vtx->fSumPt =  hartmut->ht();

  vtx->fQ       = hartmut->charge();
  vtx->fNTracks = hartmut->nTracks();
  
  const HepSymMatrix &vtxcov = hartmut->vertexError();
  const HepSymMatrix &momcov = hartmut->momentumError();
 
  vtx->fVtxCov(0,0) = vtxcov[0][0];
  vtx->fVtxCov(0,1) = vtxcov[0][1];
  vtx->fVtxCov(0,2) = vtxcov[0][2];
  vtx->fVtxCov(1,0) = vtxcov[0][1];
  vtx->fVtxCov(1,1) = vtxcov[1][1];
  vtx->fVtxCov(1,2) = vtxcov[1][2];
  vtx->fVtxCov(2,0) = vtxcov[0][2];
  vtx->fVtxCov(2,1) = vtxcov[1][2];
  vtx->fVtxCov(2,2) = vtxcov[2][2];

  vtx->fMomCov(0,0) = momcov[0][0];
  vtx->fMomCov(0,1) = momcov[0][1];
  vtx->fMomCov(0,2) = momcov[0][2];
  vtx->fMomCov(1,0) = momcov[0][1];
  vtx->fMomCov(1,1) = momcov[1][1];
  vtx->fMomCov(1,2) = momcov[1][2];
  vtx->fMomCov(2,0) = momcov[0][2];
  vtx->fMomCov(2,1) = momcov[1][2];
  vtx->fMomCov(2,2) = momcov[2][2];

  vtx->fMassError = hartmut->massError();
  vtx->fChi2      = hartmut->chi2();
  vtx->fNdof      = hartmut->ndof();
  return 0;
}

//_____________________________________________________________________________
Int_t StntupleInitVertexBlock(TStnDataBlock* block, AbsEvent* event,int mode)
{
  // initialize vertex data block with the `event' data

  int   ev_number, rn_number;
//-----------------------------------------------------------------------------
//  save time: don't do initialization twice for the same event
//-----------------------------------------------------------------------------
  ev_number = AbsEnv::instance()->trigNumber();
  rn_number = AbsEnv::instance()->runNumber();

  if (block->Initialized(ev_number,rn_number)) return 0;

  TStnVertexBlock* data = (TStnVertexBlock*) block;
  data->Clear();
					// don't read all the collections, 
					// select only the one requested

  char  process[100], description[100];
  StntupleGetProcessName(data->CollName()->Data(),process,description);
  ConstHandle<VertexColl> handle(StntupleGetIterator(event,
				  "VertexColl",process,description));

  if (handle.is_null()) {
     return -1;
  } 

  int rc = 0;
  for(VertexColl::const_iterator vertex = handle->contents().begin();
      vertex != handle->contents().end();++vertex) {
    TStnVertex* stnvtx = data->NewVertex();
    stnvtx->fNumber = data->NVertices();
    StntupleInitVertex(stnvtx,(const Vertex*)&(*(*vertex)));
  }

//-----------------------------------------------------------------------------
//  on return mark block as initialized for given event/run
//-----------------------------------------------------------------------------
  data->f_RunNumber   = rn_number;
  data->f_EventNumber = ev_number;

  return rc;
}


//_____________________________________________________________________________
// init for ZVertexModule -> ZVertexColl

Int_t StntupleInitZVertexBlock(TStnDataBlock* block, AbsEvent* event,int mode)
{
  int   ev_number, rn_number;
//-----------------------------------------------------------------------------
//  save time: don't do initialization twice for the same event
//-----------------------------------------------------------------------------
  ev_number = AbsEnv::instance()->trigNumber();
  rn_number = AbsEnv::instance()->runNumber();

  if (block->Initialized(ev_number,rn_number)) return 0;

  TStnVertexBlock* data = (TStnVertexBlock*) block;
  data->Clear();
					// don't read all the collections, 
					// select only the one requested

  std::string    coll_name = data->CollName()->Data();

  int rc = -1;
  for (EventRecord::ConstIterator it(event,"ZVertexColl","ZVertexColl");
       it.is_valid(); ++it) {
    ConstHandle<ZVertexColl> vs(it);
    rc = 0;
    for(ZVertexColl::const_iterator vertex = vs->contents().begin();
	vertex != vs->contents().end();++vertex) {
      TStnVertex* stnvtx = data->NewVertex();
      stnvtx->fNTracks = vertex->getNTracks();
      stnvtx->fVClass  = vertex->getQuality();
      stnvtx->fSumPt   = vertex->getSumPt();
      stnvtx->fT0      = vertex->getCotT0();
      stnvtx->fT0Error = vertex->getCotT0Error();
      stnvtx->fVertex  = TVector3(99.,99.,vertex->zPosition());
    }
  }

//-----------------------------------------------------------------------------
//  on return mark block as initialized for given event/run
//-----------------------------------------------------------------------------
  data->f_RunNumber   = rn_number;
  data->f_EventNumber = ev_number;

  return rc;
}