#include "PndTpcLheHitsMaker.h"

#include "PndTpcLheHit.h"
#include "PndTpcLheTrack.h"

#include "CbmRunAna.h"
#include "PndTpcPoint.h"
#include "CbmMCTrack.h"
#include "CbmMCApplication.h"
#include "CbmRootManager.h"

#include "TGeoManager.h"
#include "TGeoVolume.h"
#include "CbmVolume.h"
#include "CbmTask.h"
#include "TParticle.h"
#include "TRandom.h"
#include "Riostream.h"
#include "TH1F.h"

#include <iomanip>


//#include "lhe.h"

//#define TST
//#define TPC
#define TRACKS
#define GEANT
#define MASSPROD

//_________________________________________________________________
PndTpcLheHitsMaker::PndTpcLheHitsMaker() {
  //---

  fLheHits  = new TClonesArray("TpcLheHit");
  fGeantTracks = new TClonesArray("TpcLheTrack");
}

//_________________________________________________________________
PndTpcLheHitsMaker::PndTpcLheHitsMaker(const char *name,
				       const char *title):CbmTask(name) {
  //---
  fLheHits = new TClonesArray("TpcLheHit");
  fGeantTracks = new TClonesArray("TpcLheTrack");

}

//_________________________________________________________________
PndTpcLheHitsMaker::~PndTpcLheHitsMaker() {
  //
  CbmRootManager *fManager =CbmRootManager::Instance();
  fManager->Write();
}

//_________________________________________________________________
InitStatus PndTpcLheHitsMaker::ReInit() {
  // --- Update the pointers after reinitialisation

  //  cout << "PndTpcLheHitsMaker::ReInit() \n";

#if 0
  CbmRunAna* ana = CbmRunAna::Instance();
  CbmRuntimeDb* rtdb=ana->GetRuntimeDb();
  fGeoPar=(CbmGeoStsPar*)(rtdb->getContainer("CbmGeoStsPar"));
#endif


  return kSUCCESS;
}

//_________________________________________________________________
InitStatus PndTpcLheHitsMaker::Init() {

  //  cout << "InitStatus PndTpcLheHitsMaker::Init\n\n";

  fCuts = PndTpcLheTrackCuts::Instance();

  CbmRootManager *fManager =CbmRootManager::Instance();	

  fTpcPoints   = (TClonesArray *)fManager->GetObject("TpcPoint");
  if ( ! fTpcPoints ) {
    cout << "-I- "<< GetName() << "::Init: No TpcPoint array!" << endl;
    return kERROR;
  }

#if 0
  fTstPoints = (TClonesArray*) fManager->GetObject("TSTPoint");
  if ( ! fTstPoints ) {
    cout << "-I- "<< GetName() << "::Init: No TSTPoint array!" << endl;
    return kERROR;
  }
#endif

#ifdef GEANT
  fListMCtracks = (TClonesArray *)fManager->GetObject("MCTrack");
#endif

  fNofEvents = 0;

  Register();

  return kSUCCESS;

}

//_________________________________________________________________
void PndTpcLheHitsMaker::GetTstHits() {
  //---    

#ifdef TST
  cout << " PndTpcLheHitsMaker::GetTstHits() : TST Hit entries " <<
    fTstPoints->GetEntriesFast() << endl;
#endif

  PndTpcLheHit *hit=NULL;

#if 0
  for (int j=0; j < fTstPoints->GetEntries(); j++ ) {
    CbmTstPoint *point = (CbmTstPoint*) fTstPoints->At(j);

#ifdef KTST
    cout << " hit: " << j << " detID " <<  point->GetStationNr() <<
    " event: " << info->GetEventNumber() << " fake: " <<
      info->IsFake() << "\n ";
#endif

    hit = AddHit();
    hit->SetHitNumber(fNHit++);

    hit->SetX(point->GetX());
    hit->SetY(point->GetY());
    hit->SetZ(point->GetZ());

    hit->SetTrackID(point->GetTrackID());

    hit->SetRefIndex(j);

#if 0
    hit->SetXerr(point->GetDx());
    hit->SetYerr(point->GetDy());
    hit->SetZerr(point->GetDz());

    hit->SetStation(point->GetStationNr()); //

    if (info->GetEventNumber() == 0 ) {
      //&& !info->IsFake()) {
      //CbmMCPoint *mcpoint = (CbmMCPoint *) fTpcPoints->At(point->GetRefIndex());
      //	hit->SetTrackID(mcpoint->GetTrackID());

	hit->SetTrackID(info->GetTrackId());
	hit->SetRefIndex(point->GetRefIndex());
      }
      else {
	hit->SetTrackID(-1);
      }
#endif

#ifdef TST
    hit->Print();
#endif

  }
#endif
}

//_________________________________________________________________
void PndTpcLheHitsMaker::GetTpcHits() {
    
#ifdef TPC
  cout << " PndTpcLheHitsMaker::GetTpcHits(): Tpc points entries " <<
     fTpcPoints->GetEntriesFast() <<endl;
#endif

  for (int j=0; j < fTpcPoints->GetEntries(); j++ ) {

    PndTpcPoint* point = (PndTpcPoint*) fTpcPoints->At(j);

#ifdef TPC
    point->Print(" "); PR(point->GetTrackID());
#endif


    PndTpcLheHit* hit = AddHit();
    hit->SetHitNumber(fNHit++);

    hit->SetX(point->GetX());
    hit->SetY(point->GetY());
    hit->SetZ(point->GetZ());

    hit->SetXerr(0.03);
    hit->SetYerr(0.03);
    hit->SetZerr(0.03);

    hit->SetTrackID(point->GetTrackID());
    hit->SetRefIndex(j);
    
#ifdef TPC
    hit->Print();
#endif


  }  // for (int j=0; j <  fStsHits->GetEntries(); j++ )

}

//_________________________________________________________________
void PndTpcLheHitsMaker::Exec(Option_t * option) {
    
  cout << "\n\n  ***  Event # " << ++fNofEvents << endl;

  cout << " =====   PndTpcLheHitsMaker   =====\n";

  Reset();

  //  GetTstHits();

  GetTpcHits();

  cout << "  Total number of hits for tracking: " <<
    setw(5) << fLheHits->GetEntries() << endl;

#ifdef GEANT
  CheckTracks();
#ifdef TRACK
  PrintTracks(0);
#endif
#endif
  
#ifdef MASSPROD
  //  fTpcPoints->Clear("C");
  //  fTstPoints->Clear("C");
#endif

}

//_________________________________________________________________
PndTpcLheHit * PndTpcLheHitsMaker::AddHit() {
  // Creates a new hit in the TClonesArray.

  TClonesArray& hitRef = *fLheHits;
  Int_t size = hitRef.GetEntriesFast();
  return new(hitRef[size]) PndTpcLheHit();
}

//_________________________________________________________________
void PndTpcLheHitsMaker::Register() {
  //---

  CbmRootManager::Instance()->
    Register("LheGeantTrack","Lhe", fGeantTracks, kFALSE);
  //    Register("LheGeantTrack","Lhe", fGeantTracks, kTRUE);

  CbmRootManager::Instance()->
    Register("LheHit","Lhe", fLheHits, kFALSE);

}

//_________________________________________________________________
void PndTpcLheHitsMaker::Finish() {

//   CbmRootManager *fManager =CbmRootManager::Instance();	
//   fManager->Fill();

  cout << " PndTpcLheHitsMaker::Finish : Geant tracks " <<
    fGeantTracks->GetEntriesFast() << endl;

  fGeantTracks->Delete();
  fLheHits->Delete();

}

//_________________________________________________________________
void PndTpcLheHitsMaker::Reset() {
  //---
  //  cout << " PndTpcLheHitsMaker::Reset  " << endl;

  fNHit = 0;
  fNTrack = 0;	

  if (fLheHits->GetEntries() != 0)  fLheHits->Clear("C");
  if (fGeantTracks->GetEntries() != 0)fGeantTracks->Clear("C");

}

#ifdef GEANT
//_______________________________________________________________
void PndTpcLheHitsMaker::CheckTracks() {

  using namespace std;

  Int_t  parent, tr_num, tr_pdg;
  Int_t nhit = 0;

  TObjArray dtracks;
  TObjArray  *chits;

  CbmMCTrack *gtrack = 0;

  Int_t nMCtracks = fListMCtracks->GetEntries();

  Int_t selected_tracks = 0;

  for (Int_t ih = 0; ih < fLheHits->GetEntries(); ih++) {
    PndTpcLheHit *hit = (PndTpcLheHit *)fLheHits->UncheckedAt(ih);
    SetTrack(hit);
    //    hit->Print();
  }

  //#ifdef TRACK
  cout << " MC tracks in event: " <<fListMCtracks->GetEntries();
  cout << " tracks in TPC " << fNTrack << endl;
  //#endif

  for (Int_t itrack = 0; itrack < fNTrack; itrack++) {

    PndTpcLheTrack *track = (PndTpcLheTrack*)fGeantTracks->UncheckedAt(itrack);
    tr_num = track->GetTrackNumber();

    if(tr_num >= 0 && tr_num < nMCtracks) {
      tr_pdg = 0;
      parent = 99999;

      gtrack = (CbmMCTrack *) fListMCtracks->At(tr_num);

      tr_pdg = gtrack->GetPdgCode();
      track->SetPid(tr_pdg);
      track->SetCharge(TMath::Sign(1,tr_pdg));
      parent = gtrack->GetMotherID();
      if(parent == -1) {
	track->ComesFromMainVertex(kTRUE);
	track->SetGood(kTRUE);
      }

      TVector3 vertex = gtrack->GetStartVertex();
      track->SetVertex(vertex.X(), vertex.Y(), vertex.Z());

      TVector3 moment = gtrack->GetMomentum();
      track->SetPx(moment.X());
      track->SetPy(moment.Y());
      track->SetPz(moment.Z());

      nhit = track->GetNumberOfHits();
      chits = track->GetRHits();
      // print some info about track

      cout << "\n *** track *** " << track->GetTrackNumber() <<
	" pdg " << track->GetPid() <<
	" nhits " <<nhit << " p_t " << track->GetPt() << endl;
      
      track->Print();
      PndTpcLheHit *hit = (PndTpcLheHit *)chits->First();
      hit->Print();

      if (fCuts->IsGoodGeantTrack(track)) {
	selected_tracks++;
      }

      if (parent == 99999 ) cout << "  track parent unknown";

    }
  }

#if 0
  for (Int_t ih = 0; ih < fLheHits->GetEntries(); ih++) {
    PndTpcLheHit *hit = (PndTpcLheHit *)fLheHits->UncheckedAt(ih);
    Int_t mTrackNumber = hit->GetTrackID();
    for (Int_t itrack = 0; itrack < fNTrack; itrack++) {
      PndTpcLheTrack *track = (PndTpcLheTrack*)fGeantTracks->UncheckedAt(itrack);
      Int_t tr_num = track->GetTrackNumber();
      if (mTrackNumber ==  tr_num && !track->IsGood()) {
	hit->SetUsage(kTRUE);
      }
    }
  }
#endif

  cout << "Total number of tracks in TPC: " <<
    setw(5) << fGeantTracks->GetEntries() << endl;
  cout << "           Good tracks in TPC: " <<
    setw(5) << selected_tracks << endl;
     
}

//________________________________________________________________
void PndTpcLheHitsMaker::SetTrack(PndTpcLheHit *hit) {

  Int_t mTrackNumber = hit->GetTrackID();
  Int_t mTN;

#ifdef HITS
     cout << " fNHit = " << hit->GetHitNumber();
     cout << " fNTrack = " << fNTrack;
     cout << " mTrackNumber = " << mTrackNumber;
     cout << endl;
#endif

  PndTpcLheTrack *track = NULL;

  Bool_t newtrack = kTRUE;
  
  if (fNTrack > 0) {
    for (Int_t i = 0; i < fGeantTracks->GetEntriesFast(); i++) {
      PndTpcLheTrack *currtrk = ((PndTpcLheTrack*) fGeantTracks->At(i));
      mTN = currtrk->GetTrackNumber();
      //      PR( currtrk->GetTrackNumber());
      if (mTN == mTrackNumber) {
	track = currtrk;
	newtrack = kFALSE;
	break;
      }
    }
  }
  if (newtrack)  track = AddTrack(mTrackNumber);

  track->AddHit(hit);
  
}

//________________________________________________________________
PndTpcLheTrack *PndTpcLheHitsMaker::AddTrack(Int_t mTrackNumber) {
// Add a new track to the list of tracks for this event.

  TClonesArray &tracks = *fGeantTracks;
  PndTpcLheTrack *track = new(tracks[fNTrack++]) PndTpcLheTrack(mTrackNumber);
  return track;
}

//_________________________________________________________________
void PndTpcLheHitsMaker::PrintTracks(Int_t ntr) {


  //  cout << " " << fNTrack << " tracks  \n\n ";
  //  cout << "Station X Y Z Px Py Pz \n\n";

  Int_t  ptracks;

  if(ntr != 0 ) 
    ptracks = ntr;
  else
    ptracks = fGeantTracks->GetEntries();

  for (Int_t itrack = 0; itrack < ptracks; itrack++) {
	
    PndTpcLheTrack *track = (PndTpcLheTrack *)fGeantTracks->UncheckedAt(itrack);

    track->Print(); //

  }

}

#endif

ClassImp(PndTpcLheHitsMaker)