//*-- Author  : Rafał Lalik
//*-- Created : 01.06.2016

//_HADES_CLASS_DESCRIPTION
/////////////////////////////////////////////////////////////
//  HStsDigitizer
//
//  This class digitizes the Forward Straw detector data
//
//  Produce calibrated Time and energy loss and the straw number
//
//
//
//  !!! There digitization is not implemented yet !!!
//  Programm store to the output category catStsCal tof from geant hit
//  and minimal distance from geant track to the wire.
//
//  Geant stores three values: module, layer, cell.
//  Each module contains fixed number of layers. Each layer is composed of
/// several double-modules. Each double-module is built of single plane blocks.
//  Each block is built out of straws.
//
//  The first double-module is most left (negative x), the last is most right
//  (positive x).
//  Odd block number is for the front plane, even for the back plane. Block
//  numbers increas from left (negative x) to right (positive x).
//  Straws in a block goes from left (negative x) to right (positive x).
//
//  Cell value encodes position of single straw in a layer in a following way:
//    cell=double_module*100 + block*10 + straw
//  e.g. cell=437 -> double_module 4, block 3 (back plane), straw 7
//
//  Geometry of a single straw module for station T1(T2) is following:
//  Number of layers in a module:               4 (4)
//  Number of double-modules in a layer:        5 (7)
//  Number of blocks in a double-module:        4 (4)
//  Number of straws in a block:                8 (8)
//
//  While module and layer numbers go from 0..n, double_module, block and straw
//  numbers encoded in the cell value go 1..n
/////////////////////////////////////////////////////////////

#include "hstsdigitizer.h"
#include "hades.h"
#include "hcategory.h"
#include "hdebug.h"
#include "hevent.h"
#include "hgeantkine.h"
#include "hgeantsts.h"
#include "hgeomcompositevolume.h"
#include "hgeomvector.h"
#include "hgeomvolume.h"
#include "hruntimedb.h"
#include "hspectrometer.h"
#include "hstart2hit.h"
#include "hstartdef.h"
#include "hstscalsim.h"
#include "hstsdetector.h"
#include "hstsdigipar.h"
#include "hstsgeompar.h"
#include "stsdef.h"

#include <TRandom.h>
#include <TVector2.h>

#include <cstring>
using namespace std;

// #define VERBOSE_MODE1
// #define VERBOSE_MODE2

ClassImp(HStsDigitizer);

HStsDigitizer::HStsDigitizer()
{
    // default constructor
    initVariables();
}

HStsDigitizer::HStsDigitizer(const Text_t* name, const Text_t* title) : HReconstructor(name, title)
{
    // constructor
    initVariables();
}

void HStsDigitizer::initVariables()
{
    // Initialize the variables in constructor
    pGeantStsCat = nullptr;
    pStrawCalCat = nullptr;
    pStrawDigiPar = nullptr;
    fLoc.setNIndex(4);
    fLoc.set(4, 0, 0, 0, 0);
}

Bool_t HStsDigitizer::init()
{
    // find the Forward detector in the HADES setup
    HStsDetector* p = (HStsDetector*)(gHades->getSetup()->getDetector("Sts"));
    if (!p)
    {
        Error("StsDigitizer::init", "No Forward Detector found");
        return kFALSE;
    }

    // GEANT input data
    pGeantStsCat = gHades->getCurrentEvent()->getCategory(catStsGeantRaw);
    if (!pGeantStsCat)
    {
        Error("HStsDigitizer::init()", "HGeant Sts input missing");
        return kFALSE;
    }

    // build the Calibration category
    pStrawCalCat = p->buildCategory(catStsCal, kTRUE);
    if (!pStrawCalCat)
    {
        Error("HStsDigitizer::init()", "Cal category not created");
        return kFALSE;
    }

    // create the parameter container
    pStrawDigiPar = (HStsDigiPar*)gHades->getRuntimeDb()->getContainer("StsDigiPar");
    if (!pStrawDigiPar)
    {
        Error("HStsDigitizer::init()", "Parameter container for digitizer not created");
        return kFALSE;
    }

    // create the parameter container
    pStrawGeomPar = (HStsGeomPar*)gHades->getRuntimeDb()->getContainer("StsGeomPar");
    if (!pStrawGeomPar)
    {
        Error("HStsDigitizer::init()", "Parameter container for geometry not created");
        return kFALSE;
    }

    pStartHitCat = gHades->getCurrentEvent()->getCategory(catStart2Hit);
    if (!pStartHitCat) Warning("init", "Start hit level not defined; setting start time to 0");

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
            for (Int_t c = 0; c < STS_MAX_CELLS; ++c)
                stsCellsLab[m][l][c] = nullptr;

    return kTRUE;
}

Bool_t HStsDigitizer::reinit()
{
    dt_p[0] = pStrawDigiPar->getDriftTimePar(0);
    dt_p[1] = pStrawDigiPar->getDriftTimePar(1);
    dt_p[2] = pStrawDigiPar->getDriftTimePar(2);
    dt_p[3] = pStrawDigiPar->getDriftTimePar(3);
    dt_p[4] = pStrawDigiPar->getDriftTimePar(4);

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
        {
            HModGeomPar* fmodgeom = pStrawGeomPar->getModule(l, m);
            // If module/sector is disabled in geometry, skip it.
            if (!fmodgeom) continue;

            HGeomTransform labTrans = fmodgeom->getLabTransform();

            cosa[m][l] = labTrans.getRotMatrix().getElement(1, 1);
            sina[m][l] = -fabs(labTrans.getRotMatrix().getElement(1, 0));

            if (cosa[m][l] < 0)
            {
                cosa[m][l] = -cosa[m][l];
                sina[m][l] = -sina[m][l];
            }

            HGeomCompositeVolume* fMod = fmodgeom->getRefVolume();

            for (Int_t c = 0; c < STS_MAX_COMP; ++c)
            {
                HGeomVolume* fVol = fMod->getComponent(c);
                if (!fVol) break;
                if (stsCellsLab[m][l][c] == nullptr) stsCellsLab[m][l][c] = new HGeomVector;
                HGeomVector* p = stsCellsLab[m][l][c];
                *p = fVol->getTransform().getTransVector();
                *p = labTrans.transFrom(*p);
            }
        }
    }

    return kTRUE;
}

Int_t HStsDigitizer::execute()
{
    //--------------------------------------------------
    // start time extraction.
    // Getting start time smearing
    Float_t startTimeSmearing = 0; //[ns]
    if (pStartHitCat && pStartHitCat->getEntries() > 0)
    {
        HStart2Hit* pStartH = (HStart2Hit*)pStartHitCat->getObject(0);
        if (pStartH != NULL && pStartH->getResolution() != -1000) { startTimeSmearing = pStartH->getResolution(); }
    }

    // Digitization of GEANT hits and storage in HStsCalSim

    Char_t geaModule;  // module number (0...8); straw: 0,1
    Char_t geaLayer;   // layer number (0..8): straw: 0-4
    Int_t geaCell;     // cell number
    Char_t geaSubCell; // sub cell number

    // get all params
    Float_t adc_reso = pStrawDigiPar->getAnalogReso();
    Float_t eloss_slope = pStrawDigiPar->getElossSlope();
    Float_t eloss_offset = pStrawDigiPar->getElossOffset();
    Float_t threshold = pStrawDigiPar->getThreshold();
    Float_t efficiency = pStrawDigiPar->getEfficiency();

#ifdef VERBOSE_MODE1
    Int_t cnt = -1;
    printf("<<----------------------------------------------->>\n");
#endif

    std::map<Int_t, Int_t> track_hit_cnt;

    Int_t entries = pGeantStsCat->getEntries();
    for (Int_t i = 0; i < entries; ++i)
    {
        HGeantSts* ghit = (HGeantSts*)pGeantStsCat->getObject(i);
        if (ghit)
        {
#ifdef VERBOSE_MODE1
            ++cnt;
#endif
            ghit->getAddress(geaModule, geaLayer, geaCell, geaSubCell);

            // If sector is disabled in setup, skip it.
            if (!stsCellsLab[(int)geaModule][(int)geaLayer][0]) continue;

            // detection efficiency
            Float_t rnd = gRandom->Rndm();
#ifdef VERBOSE_MODE1
            printf("(%2d) det eff: rand=%.2f  <?<  eff=%.2f\n", cnt, rnd, efficiency);
#endif
            if (rnd > efficiency) continue;

            DigiFields df;
            GeantFields gf;

            df.m = geaModule;
            df.l = geaLayer;

            // Straw number calculation
            // Straw and offsets are organized as follow (l: long, s: short)
            //
            // llllllllssssllllllll
            // llllllllssssllllllll
            // llllllllssssllllllll
            // llllllllssssllllllll
            // llllllll    llllllll
            // llllllll    llllllll
            // llllllll    llllllll
            // llllllll    llllllll
            // llllllllssssllllllll
            // llllllllssssllllllll
            // llllllllssssllllllll
            // llllllllssssllllllll
            // |
            // |- offset 0 (o0)
            //         |- offset 1 (o1)
            //             |- offset 2 (o2)
            //                    |- offset 3 (03)
            //
            // If abs(straw_number) < 256: long straws, otherwsie short
            // For long:
            //   if straw_number >= o1: straw_number + o2;
            // For short: abs(straw_number) - 256 + o1

            df.s = geaCell;
            df.ud = geaSubCell;

            // straws are rotated in the geo by 90 degrees around x axis
            // y and z must be swap
            ghit->getHit(gf.xHit, gf.zHit, gf.yHit, gf.pxHit, gf.pzHit, gf.pyHit, gf.tofHit, gf.trackLength, gf.eHit);
            gf.trackNumber = ghit->getTrackNumber();

            // in the sim, z-axis is inverted
            gf.pzHit = -gf.pzHit;
            gf.zHit = -gf.zHit;

            gf.lab_x = stsCellsLab[df.m][df.l][df.s]->X() - gf.xHit * sina[df.m][df.l];
            gf.lab_y = stsCellsLab[df.m][df.l][df.s]->Y() + gf.yHit * cosa[df.m][df.l];
            gf.lab_z = stsCellsLab[df.m][df.l][df.s]->Z() + gf.zHit;

            //             TVector2 v_n(gf.pzHit, gf.pxHit);   // track vector dir.  n
            //             v_n = v_n/v_n.Mod();
            //             TVector2 p_a(hit_z, hit_x);         // track hit position a
            //             TVector2 p_p(cell_z, cell_x);       // wire position      p

            // calculations:
            // dist = | (a-p) - ((a-p)*n)n|
            //             TVector2 a_p_diff = p_a - p_p;
            //             TVector2 v_proj_n_diff = (a_p_diff * v_n) * v_n;
            //             TVector2 v_dist = a_p_diff - v_proj_n_diff;
            df.radius = abs(gf.xHit); // v_dist.Mod();

            df.time = gf.tofHit + calcDriftTime(df.radius);
            df.adc = gf.eHit * eloss_slope + eloss_offset;

#ifdef VERBOSE_MODE1
            printf("     g (m=%d l=%d c=%03d ud=%d, tr=%d) -> el=%.4f tof=%.3f "
                   "dr=%f  adc=%.2f time=%f "
                   " p=%f (%f,%f) loc=(%f,%f,%f)\n",
                   geaModule, geaLayer, geaCell, geaSubCell, gf.trackNumber, gf.eHit, gf.tofHit,
                   df.radius, df.adc, df.time, sqrt(gf.pxHit * gf.pxHit + gf.pyHit * gf.pyHit + gf.pzHit * gf.pzHit),
                   gf.pxHit / gf.pzHit, gf.pyHit / gf.pzHit, gf.lab_x, gf.lab_y, gf.lab_z);
#endif

            if (adc_reso > 0.0)
            {
                df.adc += +gRandom->Gaus() * adc_reso;
                df.adc = TMath::Max(Float_t(0), df.adc);
            }

            if (df.adc < threshold) continue;

            df.time = gRandom->Gaus(df.time, pStrawDigiPar->getDriftTimeReso()) - startTimeSmearing;

#ifdef VERBOSE_MODE1
            printf("     resolution effects                                                        "
                   "     adc=%.2f time=%f\n",
                   df.adc, df.time);
#endif

            df.posU = stsCellsLab[df.m][df.l][df.s]->X() * cosa[df.m][df.l] +
                      stsCellsLab[df.m][df.l][df.s]->Y() * sina[df.m][df.l];
            df.posZ = stsCellsLab[df.m][df.l][df.s]->Z();

            df.hitnr = ++track_hit_cnt[gf.trackNumber];

            Int_t res = fillStrawCalSim(df, gf);
            if (res == 1)
            {
                Error("HStsDigitizer::execute()", "Can't fill from %d for m=%d l=%d s=%d ud=%d", geaCell, df.m, df.l,
                      df.s, df.ud);
                //                 return kFALSE;
            }
            //             if (res == 2)
            //             {
            //                 Error("HStsDigitizer::execute()",
            //                 "Overwrite fill from %d for m=%d l=%d p=%d s=%d ud=%d", geaCell,
            //                 df.mod, df.lay, df.plane, df.straw, df.udconf);
            // //                 return kFALSE;
            //             }
        }
    }

#ifdef VERBOSE_MODE1
    printf("<<- done ---------------------------------------->>\n");
#endif

    return 0;
}

Int_t HStsDigitizer::fillStrawCalSim(const DigiFields& df, const GeantFields& gf)
{
    // Function creat and fill object HStsCalSim
    // return:
    //  0 - if stored properly
    //  1 - if data was not stored to the category
    //  2 - if multiple hit, and data were not owerwritten
    fLoc[0] = df.m;
    fLoc[1] = df.l;
    fLoc[2] = df.s;
    fLoc[3] = df.ud;

    Int_t first = 1;

    HStsCalSim* cal = (HStsCalSim*)pStrawCalCat->getObject(fLoc);
    if (cal != nullptr) // straw tube ocuppied by another track
    {
        if (cal->getTime() < df.time)
        {
            //             cal->setTrack(gf.trackNumber);   // FIXME should we have multiple track
            //             numbers?
            return 2;
        }
        first = 0;
    }
    else { cal = (HStsCalSim*)pStrawCalCat->getSlot(fLoc); }
    if (cal)
    {
        if (first) cal = new (cal) HStsCalSim;
        cal->setAddress(fLoc[0], fLoc[1], fLoc[2], fLoc[3]);

        cal->setHit(df.time, df.adc, df.posU, df.posZ);
        cal->setDriftRadius(df.radius);

        cal->setTof(gf.tofHit);
        cal->setEloss(gf.eHit);
        cal->setTrack(gf.trackNumber);
        cal->setP(gf.pxHit, gf.pyHit, gf.pzHit);
        cal->setHitPos(gf.lab_x, gf.lab_y, gf.lab_z);

        cal->setHitNumber(df.hitnr);
    }
    else
        return 1;

    return 0;
}

Float_t HStsDigitizer::calcDriftTime(Float_t x) const
{
    // Calculate dirft time using simulated drift radius
    Float_t x2 = x * x;
    Float_t x3 = x2 * x;
    Float_t x4 = x3 * x;

    return dt_p[0] + dt_p[1] * x + dt_p[2] * x2 + dt_p[3] * x3 + dt_p[4] * x4;
}