//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Task to convert ShieldPoints to CdcTracks
//
// Environment:
//      Software not developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sverre DoerheimTUM       (original author)
//-----------------------------------------------------------

#include <iostream>
#include <set>

#include <TRandom.h>

#include <FairRun.h>
#include <FairRuntimeDb.h>
#include <PndMCTrack.h>

#include <GFTrack.h>
#include <RKTrackRep.h>

#include "CdcHit.h"
#include "TpcShieldPoint.h"
#include "TpcShieldToCdcTrackTask.h"

TpcShieldToCdcTrackTask::TpcShieldToCdcTrackTask()
    : fpersistence(false), fPointBranchName("TpcShieldPoint"),
      fMCTrackBranchName("MCTrack"), fTrackOutBranchName("CdcGFTrack"),
      fCdcTrackOutBranchName("CdcTrack"), fPointArray(NULL),
      fTrackOutArray(NULL), fCdcTrackOutArray(NULL), fMCTrackArray(NULL),
      fCdcHitOutArray(NULL),

      fPhiSmearing(5 *TMath::DegToRad()), fThetaSmearing(10 *TMath::RadToDeg()),
      fMomentumSmearing(7/100), fPosSmearing(0.2, 0.2, 5), fB(0.616),
      fpar(NULL) {}

TpcShieldToCdcTrackTask::~TpcShieldToCdcTrackTask() {}

void TpcShieldToCdcTrackTask::SetParContainers() {

  std::cout << "TpcShieldToCdcTrackTask::SetParContainers" << std::endl;
  std::cout.flush();

  // Get run and runtime database
  FairRun *run = FairRun::Instance();
  if (!run)
    Fatal("SetParContainers", "No analysis run");

  FairRuntimeDb *db = run->GetRuntimeDb();
  if (!db)
    Fatal("SetParContainers", "No runtime database");

  // Get Tpc digitisation parameter container
  fpar = (TpcDigiPar *)db->getContainer("TpcDigiPar");
  if (!fpar)
    Fatal("SetParContainers", "TpcDigiPar not found");
}

InitStatus TpcShieldToCdcTrackTask::Init() {
  // Get ROOT Manager
  FairRootManager *ioman = FairRootManager::Instance();

  if (ioman == 0) {
    Error("TpcClusterizerTask::Init", "RootManager not instantiated!");
    return kERROR;
  }

  // Get input collection
  fPointArray = (TClonesArray *)ioman->GetObject(fPointBranchName);

  if (fPointArray == 0) {
    Error("TpcClusterizerTask::Init", "Point-array not found!");
    throw;
    return kERROR;
  }
  fMCTrackArray = (TClonesArray *)ioman->GetObject(fMCTrackBranchName);
  if (fMCTrackArray == 0) {
    Error("TpcClusterizerTask::Init", "MC track -array not found!");
    throw;
    return kERROR;
  }
  // create and register output array
  fTrackOutArray = new TClonesArray("GFTrack");
  ioman->Register(fTrackOutBranchName, "Tpc", fTrackOutArray, fpersistence);
  fCdcTrackOutArray = new TClonesArray("CdcTrack");
  ioman->Register(fCdcTrackOutBranchName, "Tpc", fCdcTrackOutArray,
                  fpersistence);
  fCdcHitOutArray = new TClonesArray("CdcHit");
  ioman->Register("CdcHit", "Tpc", fCdcHitOutArray, fpersistence);

  return kSUCCESS;
}

void TpcShieldToCdcTrackTask::Exec(Option_t *opt) {

  fTrackOutArray->Delete();
  fCdcHitOutArray->Delete();
  fCdcTrackOutArray->Delete();
  int nPoints = fPointArray->GetEntriesFast();
  std::set<int> trackIds;
  int nMCTracks = fMCTrackArray->GetEntriesFast();
  for (int iPoint = 0; iPoint < nPoints; ++iPoint) {

    TpcShieldPoint *point = (TpcShieldPoint *)fPointArray->UncheckedAt(iPoint);
    int trackId = point->GetTrackID();

    if (trackIds.count(trackId) == 0) {
      trackIds.insert(trackId);
      TVector3 pos, mom, poserr, momerr;
      point->Momentum(mom);
      point->Position(pos);
      TVector3 origPos = pos;
      TVector3 blaPos = pos;
      if (pos.Z() > 50) {
        continue;
      }

      TVector3 posSpread(fPosSmearing.X(), fPosSmearing.Y(), fPosSmearing.Z());
      double momSpread = mom.Mag() * fMomentumSmearing;
      double phiSpread = fPhiSmearing;
      double thetaSpread = fThetaSmearing;

      poserr.SetXYZ(0.2, 0.2, 10.);
      momerr.SetXYZ(mom.x() * 0.05, mom.y() * 0.05, mom.z() * 0.05);
      mom.SetMagThetaPhi(gRandom->Gaus(mom.Mag(), momSpread),
                         gRandom->Gaus(mom.Theta(), thetaSpread),
                         gRandom->Gaus(mom.Phi(), phiSpread));
      pos.SetXYZ(gRandom->Gaus(pos.X(), posSpread.X()),
                 gRandom->Gaus(pos.Y(), posSpread.Y()),
                 gRandom->Gaus(pos.Z(), posSpread.Z()));
      int pdg = 211;
      if (trackId >= nMCTracks) {
        std::cout << "trackId out of range, secId: " << point->GetSecID()
                  << std::endl;
        continue;
      }
      PndMCTrack *mcTrack =
          ((PndMCTrack *)(fMCTrackArray->At(point->GetTrackID())));
      pdg = mcTrack->GetPdgCode();
      std::cout << "adding track with: pdg: " << pdg << ", mom: " << mom.Mag()
                << " " << mcTrack->GetMomentum().Mag() << " posZ " << pos.Z()
                << std::endl;
      RKTrackRep *rkrep = new RKTrackRep(pos, mom, poserr, momerr, pdg);
      // store GFTracks in output array
      GFTrack *gftrk = new ((*fTrackOutArray)[fTrackOutArray->GetEntriesFast()])
          GFTrack(rkrep);
      CdcTrack *cdctrk =
          new ((*fCdcTrackOutArray)[fCdcTrackOutArray->GetEntriesFast()])
          CdcTrack(createCdcTrack(mom, pos, TMath::Sign(1, pdg)));
      
      cdctrk->SetTrackID(fCdcTrackOutArray->GetEntriesFast() - 1);
      origPos.SetXYZ(gRandom->Gaus(blaPos.x(), posSpread.X() ),
                     gRandom->Gaus(blaPos.y(), posSpread.Y() ),
                     gRandom->Gaus(blaPos.z(), posSpread.Z()));
      CdcHit *cdcHit =
          new ((*fCdcHitOutArray)[fCdcHitOutArray->GetEntries()]) CdcHit();
      cdcHit->SetHitPos(origPos.x(), origPos.y(), origPos.z());
      std::cout << "hitPos: ";
      origPos.Print();

      cdcHit->SetHitID(fCdcHitOutArray->GetEntries() - 1);
      cdctrk->addHit(fCdcHitOutArray->GetEntries() - 1);

      // create cdc track here
    }
  }
}

CdcTrack TpcShieldToCdcTrackTask::createCdcTrack(TVector3 mom, TVector3 pos,
                                                 int charge) {
  double r = mom.Perp() / (fB * 0.3) * 100;

  TVector3 radvec = mom;
  radvec.SetZ(0);
  radvec.RotateZ(charge * TMath::PiOver2());
  radvec.SetMag(r);
  TVector3 center = pos + radvec;
  double theta = mom.Theta();
  CdcTrack track;
  track.SetTrackID(0);
  track.SetNpoint(60);
  track.SetMom(mom.Mag());
  track.SetCharge(charge);
  track.SetRadius(r);
  track.SetBMass(0);
  track.SetMass(0);
  track.SetTheta(theta);
  track.SetPhi(0);
  track.SetEloss(0);
  track.SetMx(center.x());
  track.SetMy(center.y());
  return track;
}