// -------------------------------------------------------------------------
// -----       CbmSttVertexFinderIdeal source file      -----
// -----                  Created 28/03/06  by V. Friese               -----
// -------------------------------------------------------------------------

#include "CbmSttVertexFinderIdeal.h"

#include "CbmRootManager.h"
#include "CbmMCTrack.h"
#include "CbmSttVertex.h"
#include "CbmSttTrackMatch.h"

#include <vector>

using std::vector;

// -----   Default constructor   -------------------------------------------
CbmSttVertexFinderIdeal::CbmSttVertexFinderIdeal() 
{ 
  fMCTrackArray = NULL;
  fVerbose      = 0;
}
// -------------------------------------------------------------------------

// -----   Standard constructor   ------------------------------------------
CbmSttVertexFinderIdeal::CbmSttVertexFinderIdeal(Int_t verbose) 
{ 
  fMCTrackArray = NULL;
  fVerbose      = verbose;
}
// -------------------------------------------------------------------------

// -----   Destructor   ----------------------------------------------------
CbmSttVertexFinderIdeal::~CbmSttVertexFinderIdeal() 
{ 
}
// -------------------------------------------------------------------------

// -----   Public method Init   --------------------------------------------
void CbmSttVertexFinderIdeal::Init() 
{
  // Get and check CbmRootManager
  CbmRootManager* ioman = CbmRootManager::Instance();

  if (!ioman) 
    {
      cout << "-E- CbmSttVertexFinderIdeal::Init: "
	   << "RootManager not instantised!" << endl;
      return;
    }

  // Get MCTrack array
  fMCTrackArray  = (TClonesArray*) ioman->ActivateBranch("MCTrack");
  if ( ! fMCTrackArray) 
    {
      cout << "-E- CbmSttVertexFinderIdeal::Init: No MCTrack array!"
	   << endl;
      return;
    }

  // Get SttTrackMatch array
  fMatches = (TClonesArray*) ioman->GetObject("STTTrackMatch");
  if ( ! fMatches ) {
    cout << "-E- " << GetName() << "::Init: No SttTrackMatch array!" 
	 << endl;
    return;
  }

  
}
// -------------------------------------------------------------------------

// -----   Public method DoFind   ------------------------------------------
Int_t CbmSttVertexFinderIdeal::DoFind(TClonesArray* vertexArray) 
{
  Int_t
    numberOfVertices = 0;

  // Check pointers
  if ( !fMCTrackArray ) 
    {
      cout << "-E- CbmSttVertexFinderIdeal::DoFind: "
	   << "MCTrack array missing! " << endl;
      return -1;
    }
    
  if ( !vertexArray ) 
    {
      cout << "-E- CbmSttVertexFinderIdeal::DoFind: "
	   << "Vertex array missing! " << endl;
      return -1;
    }

  // create map
  map <int, vector<int> >
    decayMap;

  Int_t nMCTracks = fMCTrackArray->GetEntriesFast();

  // loop over all MC particles 
  for (Int_t iMCTrack = 0; iMCTrack < nMCTracks; iMCTrack++) 
  {
    CbmMCTrack*
      pMCtr = (CbmMCTrack*) fMCTrackArray->At(iMCTrack);

    if ( ! pMCtr ) 
      continue;

    // enter into map primaries have mother = -1
    if (pMCtr->GetMotherID() + 1)
      decayMap[pMCtr->GetMotherID()].push_back(iMCTrack);
  }

  map <int, vector<int> >::iterator
    decayMapIter = decayMap.begin();

  // loop over map
  while (decayMapIter != decayMap.end())
    {
      // if there was a 2 particle decay
      if (decayMapIter->second.size() == 2)
	{
	  Int_t
	    recoTrack1 = IsMCTrackFound(decayMapIter->second[0]),
	    recoTrack2 = IsMCTrackFound(decayMapIter->second[1]);

	  if (fVerbose)
	    {
	      CbmMCTrack*
		pMCtr = (CbmMCTrack*) fMCTrackArray->At(decayMapIter->first);

	      cout << "2 body decay: " << pMCtr->GetPdgCode() << " decays into: " << endl;
	      
	      pMCtr = (CbmMCTrack*) fMCTrackArray->At(decayMapIter->second[0]);
	      cout << "     "  << pMCtr->GetPdgCode() << endl;
	      pMCtr = (CbmMCTrack*) fMCTrackArray->At(decayMapIter->second[1]);
	      cout << "     "  << pMCtr->GetPdgCode() << endl;
	    }
 
	  // check if tracks have been found in tracker
	  if ((recoTrack1 >= 0) && (recoTrack2 >= 0))
	    {
	      // create new vertex object
	      CbmSttVertex
		*myVertex = new((*vertexArray)[numberOfVertices]) CbmSttVertex();
	      
	      numberOfVertices++;

	      cout << "storing track in vertex: " << recoTrack1 << endl;
	      cout << "storing track in vertex: " << recoTrack2 << endl;

	      myVertex->AddTrack(recoTrack1, decayMapIter->second[0]); // reonstructed trackno, mc trackno
	      myVertex->AddTrack(recoTrack2, decayMapIter->second[1]);
	    }
	}

      decayMapIter++;
    }
  
  return 0;
}

Int_t CbmSttVertexFinderIdeal::IsMCTrackFound(Int_t trackID)
{
  Int_t 
    nMatch = fMatches->GetEntriesFast();

  for (Int_t iMatch = 0; iMatch < nMatch; iMatch++) 
    {
      CbmSttTrackMatch
	*match = (CbmSttTrackMatch*) fMatches->At(iMatch);

      if (match->GetMCTrackID() == trackID)
	return iMatch;
    }

  return -1;
}

ClassImp(CbmSttVertexFinderIdeal)