/////////////////////////////////////////////////////////////////////////////
//
//  fill FwDet branches of GEANT ROOT tree
//
// last modified on 30/11/2015 by I. Koenig
/////////////////////////////////////////////////////////////////////////////

#include "fillfwdet.h"
#include "cfortran.h"
#include "hgeantfwdet.h"
#include "hgeantmaxtrk.h"
#include "hades.h"
#include "hlocation.h"
#include "hcategory.h"
#include "hrecevent.h"
#include "hpartialevent.h"
#include "hgeantkine.h"

#include <map>

struct FWDET_DEF {    // structure mapped to FWDETTUPLE Fortran COMMON block
  Int_t   ntrk;
  Int_t   fwdettrk[MAXTRKFWDET];
  Int_t   fwdetpart[MAXTRKFWDET];
  Int_t   fwdetmod[MAXTRKFWDET];
  Int_t   fwdetlayer[MAXTRKFWDET];
  Int_t   fwdetcell[MAXTRKFWDET];
  Int_t   fwdetsub[MAXTRKFWDET];
  Float_t fwdetx[MAXTRKFWDET];
  Float_t fwdety[MAXTRKFWDET];
  Float_t fwdetz[MAXTRKFWDET];
  Float_t fwdetpx[MAXTRKFWDET];
  Float_t fwdetpy[MAXTRKFWDET];
  Float_t fwdetpz[MAXTRKFWDET];
  Float_t fwdettof[MAXTRKFWDET];
  Float_t fwdetlen[MAXTRKFWDET];
  Float_t fwdete[MAXTRKFWDET];
};

#define FWDETTUPLE COMMON_BLOCK(FWDETTUPLE,fwdettuple)
COMMON_BLOCK_DEF(FWDET_DEF,FWDETTUPLE);

typedef std::map<Int_t,Int_t>::const_iterator ItTrk;
typedef std::multimap<Int_t,Hit>::const_iterator ItHit;

// http://stackoverflow.com/questions/1903954/is-there-a-standard-sign-function-signum-sgn-in-c-c
template <typename T> Int_t sgn(T val)
{
    return (T(0) < val) - (val < T(0));
}

bool same_hits(const Hit & h1, const Hit & h2)
{
    return (h1.track == h2.track) && (h1.module == h2.module) &&
             (h1.layer == h2.layer) && (h1.cell == h2.cell) &&
             (h1.subcell == h2.subcell);
}

void fillfwdet(void)
{
   HRecEvent* pHGeantEvent = (HRecEvent*)(gHades->getCurrentEvent());
                                                        // get pointer to event
   HPartialEvent* pSimul = pHGeantEvent->getPartialEvent(catSimul);
   HCategory* pFwDetCat = pSimul->getCategory(catFwDetGeantRaw);  // get category

   pFwDetCat->Clear();   // clear category if 1st sector

   HCategory* pKin = pSimul->getCategory(catGeantKine); // get KINE cat

   Int_t gmodule, glayer, gcell, gsub;

   Hit old_hit;
   old_hit.track = -1;
   old_hit.module = -1;
   old_hit.layer = -1;
   old_hit.cell = -1;
   old_hit.subcell = -1;

   std::map<Int_t,Int_t> rpc_tracks;
   std::multimap<Int_t,Hit> rpc_hits;

   // loop over tracks (TClonesArray, expands if necessary)
   for(Int_t idx = 0; idx < FWDETTUPLE.ntrk; ++idx)
   {
      gmodule = FWDETTUPLE.fwdetmod[idx]-1;  //GEANT indexing is from 1
      glayer = FWDETTUPLE.fwdetlayer[idx]-1;
      gcell = FWDETTUPLE.fwdetcell[idx]-1;
      gsub = FWDETTUPLE.fwdetsub[idx];

      Hit hit;
      hit.module = -1;
      hit.layer = -1;
      hit.cell = -1;
      hit.subcell = -1;

      hit.index = idx;
      hit.track = FWDETTUPLE.fwdettrk[idx];
      hit.px = FWDETTUPLE.fwdetpx[idx];
      hit.py = FWDETTUPLE.fwdetpy[idx];
      hit.pz = FWDETTUPLE.fwdetpz[idx];

      hit.x = FWDETTUPLE.fwdetx[idx];
      hit.y = FWDETTUPLE.fwdety[idx];
      hit.z = FWDETTUPLE.fwdetz[idx];

      hit.tof = FWDETTUPLE.fwdettof[idx];
      hit.trklen = FWDETTUPLE.fwdetlen[idx];
      hit.eloss = FWDETTUPLE.fwdete[idx];

      // straws
      if (gmodule < 2)
      {
         hit.module = gmodule;
         hit.layer = glayer;
         hit.cell = gcell;
         hit.subcell = gsub;

         // if this is reentry of the track to the same volume, ignore it
         if (same_hits(hit, old_hit))
            continue;

         fillHit(pKin, pFwDetCat, hit);

         old_hit.module = gmodule;
         old_hit.layer = glayer;
         old_hit.cell = gcell;
         old_hit.subcell = gsub;
         old_hit.track = hit.track;
      }
      // RPC
      else if (gmodule < 7)
      {
         hit.module = gmodule;
         hit.layer = glayer;
         hit.cell = gcell;
         hit.subcell = gsub-1;

         // store tracks numbers
         ++rpc_tracks[hit.track];
         //store hits assigned to the track
         rpc_hits.insert(std::pair<Int_t,Hit>(hit.track, hit));
      }
   }

   // here make rpc stuff
   ItTrk itt;
   ItHit ith;

   // these variables keep tracking of track direction
   Int_t dir_x = 0; // x,y,z
   Int_t dir_y = 0;
   Int_t dir_z = 0;
   Int_t dir_l = -1; // layer

   // used to recalculate average values for a track
   Float_t avg_x = 0.0, avg_y = 0.0, avg_z = 0.0;
   Float_t avg_px = 0.0, avg_py = 0.0, avg_pz = 0.0;
   Float_t avg_tof = 0.0, avg_trklen = 0.0;
   Int_t avg_index = 0;

   // keep array of active gaps
   UInt_t front_gap_mask = 0, rear_gap_mask = 0;

   // iterate over all track ids
   for (itt = rpc_tracks.begin(); itt != rpc_tracks.end(); ++itt)
   {
      Int_t track = itt->first;

      // get range of iterators for given key(track)
      std::pair<ItHit,ItHit> hiters = rpc_hits.equal_range(track);

      // hit will have the same module number always
      Hit reco_hit;
      reco_hit.track = track;
      reco_hit.module = hiters.first->second.module;

      Int_t gaps_count = 0; // count gaps number in a track
      Int_t sub_track = 0;  // count number of track splits
                            // (e.g. changed direction, different module)

      // loop over all hits from given track
      for (ItHit ith = hiters.first; ith != hiters.second; ++ith)
      {
         Hit h = ith->second;   // hit info

         // obtain direction values for selected hit
         Int_t curr_dir_x = sgn(h.px);
         Int_t curr_dir_y = sgn(h.py);
         Int_t curr_dir_z = sgn(h.pz);
         Int_t curr_dir_l = h.layer;

         // it this is the first hit to follow
         if (dir_l == -1)
         {
            // define direction
            dir_x = curr_dir_x;
            dir_y = curr_dir_y;
            dir_z = curr_dir_z;
            dir_l = curr_dir_l;

            avg_index = h.index;
         }
         else
         {
            // if track has changed direction, make insertion of the current track
            if (((dir_z != curr_dir_z) && (dir_x != curr_dir_x || dir_y != curr_dir_y)) || (dir_l != curr_dir_l))
            {
               reco_hit.layer = dir_l;
               reco_hit.cell = (front_gap_mask << 8) + rear_gap_mask;
               reco_hit.subcell = sub_track;

               reco_hit.x = avg_x/gaps_count;
               reco_hit.y = avg_y/gaps_count;
               reco_hit.z = avg_z/gaps_count;
               reco_hit.px = avg_px/gaps_count;
               reco_hit.py = avg_py/gaps_count;
               reco_hit.pz = avg_pz/gaps_count;
               reco_hit.tof = avg_tof/gaps_count;
               reco_hit.trklen = avg_trklen/gaps_count;
               reco_hit.index = avg_index;

               fillHit(pKin, pFwDetCat, reco_hit);

               // new direction
               dir_x = curr_dir_x;
               dir_y = curr_dir_y;
               dir_z = curr_dir_z;
               dir_l = curr_dir_l;

               avg_x = 0.0; avg_y = 0.0; avg_z = 0.0;
               avg_px = 0.0; avg_py = 0.0; avg_pz = 0.0;
               avg_tof = 0.0; avg_trklen = 0.0;
               front_gap_mask = 0; rear_gap_mask = 0;
               gaps_count = 0;

               avg_index = h.index;
               ++sub_track;
            }
         }

         avg_x += h.x;
         avg_y += h.y;
         avg_z += h.z;
         avg_px += h.px;
         avg_py += h.py;
         avg_pz += h.pz;
         avg_tof += h.tof;
         avg_trklen += h.trklen;

         if (h.cell == 0)
             front_gap_mask |= (1 << h.subcell);
         else
             rear_gap_mask |= (1 << h.subcell);

         ++gaps_count;
      }

      // after the last track, fill it as well
      reco_hit.layer = dir_l;
      reco_hit.cell = (front_gap_mask << 8) + rear_gap_mask;
      reco_hit.subcell = sub_track;

      reco_hit.x = avg_x/gaps_count;
      reco_hit.y = avg_y/gaps_count;
      reco_hit.z = avg_z/gaps_count;
      reco_hit.px = avg_px/gaps_count;
      reco_hit.py = avg_py/gaps_count;
      reco_hit.pz = avg_pz/gaps_count;
      reco_hit.tof = avg_tof/gaps_count;
      reco_hit.trklen = avg_trklen/gaps_count;
      reco_hit.index = avg_index;

      fillHit(pKin, pFwDetCat, reco_hit);
      dir_l = -1;
      avg_x = 0.0; avg_y = 0.0; avg_z = 0.0;
      avg_px = 0.0; avg_py = 0.0; avg_pz = 0.0;
      avg_tof = 0.0; avg_trklen = 0.0;
      front_gap_mask = 0; rear_gap_mask = 0;
      gaps_count = 0;
   }
}

Bool_t fillHit(HCategory* pKin, HCategory* pFwDetCat, const Hit & hit)
{
   HLocation loc;
   loc.set(0);

   HGeantFwDet* pFwDet = (HGeantFwDet*)pFwDetCat->getNewSlot(loc);  // get slot

   if (pFwDet == NULL)
   {
      printf("\n*** FILLFWDET: HGeantFwDet location not found ! ***\n");
      printf("index=%d\n", hit.index);
      return kFALSE;
   }

   pFwDet = new(pFwDet) HGeantFwDet;   // FwDet track data

   if (!pFwDet)
       return kFALSE;

   pFwDet->setTrack(hit.track);
   pFwDet->setAddress(hit.module, hit.layer, hit.cell, hit.subcell);
   pFwDet->setHit(hit.x, hit.y, hit.z, hit.px, hit.py, hit.pz, hit.tof, hit.trklen, hit.eloss);

   // add hit to linked list of track
   if (pKin != NULL && hit.track > 0)
   {
      ((HGeantKine*)pKin->getObject(hit.track-1))->setFwDetHitIndex(hit.index);
   }

   return kTRUE;
}

// wrapping of the c++ routine as Fortran-subroutine using cfortran.h of Cern-library
FCALLSCSUB0(fillfwdet,FILLFWDET,fillfwdet)