//*-- AUTHOR : A.Zinchenko <Alexander.Zinchenko@jinr.ru>
//*-- Created : 08/08/2016
//*-- Modified: R. Lalik <Rafal.Lalik@ph.tum.de> 2016/10/10

//_HADES_CLASS_DESCRIPTION
/////////////////////////////////////////////////////////////
//
//  HForwardCandFinder
//
//  Tracking code for Sts
//
/////////////////////////////////////////////////////////////

#include "hforwardbasictracker.h"
#include "forwarddef.h"
#include "frpcdef.h"
#include "hades.h"
#include "hcategory.h"
#include "hcategorymanager.h"
#include "hevent.h"
#include "hforwardcand.h"
#include "hforwardcandfinderpar.h"
#include "hforwardcandsim.h"
#include "hforwardtools.h"
#include "hfrpcdetector.h"
#include "hfrpchit.h"
#include "hgeantkine.h"
#include "hgeomcompositevolume.h"
#include "hgeominterface.h"
#include "hgeomrootbuilder.h"
#include "hgeomvector.h"
#include "hgeomvolume.h"
#include "hlinearcategory.h"
#include "hruntimedb.h"
#include "hspectrometer.h"
#include "hstscalpar.h"
#include "hstscalsim.h"
#include "hstsdetector.h"
#include "hstsdigipar.h"
#include "hstsgeompar.h"

#include <TGeoBBox.h>
#include <TGeoManager.h>
#include <TMath.h>
#include <TMatrixD.h>
#include <TMatrixFLazy.h>

#include <TKDTree.h>

#include <cassert>
#include <iterator>

// #define VERBOSE_MODE
// #define VERBOSE_MODE_PAIRS
// #define VERBOSE_MODE_CHKPARS
// #define VERBOSE_MODE_LRT
// #define VERBOSE_MODE4

HForwardBasicTracker::HForwardBasicTracker(const Text_t* name, const Text_t* title)
    : HReconstructor(name, title), pStrawCal(nullptr), pForwardCand(nullptr), pStrawVFPar(nullptr),
      isSimulation(kFALSE), n_events(0), n_ok_events(0), n_vecs(0)
{
}

HForwardBasicTracker::~HForwardBasicTracker()
{
    clear();

    for (map<Int_t, TDecompLU*>::iterator it = fLus.begin(); it != fLus.end(); ++it)
    {
        delete it->second;
    }

    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)
            {
                delete stsCellsLab[m][l][c];
                delete stsCellsLoc[m][l][c];
            }
}

Bool_t HForwardBasicTracker::init()
{
    HStsDetector* pSts = (HStsDetector*)(gHades->getSetup()->getDetector("Sts"));
    if (!pSts)
    {
        Error("HForwardBasicTracker::init", "No STS detector found");
        return kFALSE;
    }

    HFRpcDetector* pFRpc = (HFRpcDetector*)(gHades->getSetup()->getDetector("FRpc"));
    if (!pFRpc)
    {
        Error("HForwardBasicTracker::init", "No Forward RPC found");
        return kFALSE;
    }

    // GEANT input data
    pKine = gHades->getCurrentEvent()->getCategory(catGeantKine);
    if (pKine) { isSimulation = kTRUE; }

    // straw hits
    pStrawCal = gHades->getCurrentEvent()->getCategory(catStsCal);
    if (!pStrawCal)
    {
        Error("HForwardBasicTracker::init()", "HStsCal input missing");
        return kFALSE;
    }

    // build the Straw vector category
    if (isSimulation)
        pForwardCand = new HLinearCategory("HForwardCandSim");
    else
        pForwardCand = new HLinearCategory("HForwardCand");

    if (!pForwardCand)
    {
        Error("HForwardBasicTracker::init()", "ForwardCand category not created");
        return kFALSE;
    }
    else { gHades->getCurrentEvent()->addCategory(catForwardCand, pForwardCand, "Forward"); }

    pStrawVFPar = (HForwardCandFinderPar*)gHades->getRuntimeDb()->getContainer("ForwardCandFinderPar");
    if (!pStrawVFPar)
    {
        Error("HForwardBasicTracker::init()", "Parameter container for vector finder parameters not created");
        return kFALSE;
    }

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

    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;
                stsCellsLoc[m][l][c] = nullptr;
            }

    pCalPar = (HStsCalPar*)gHades->getRuntimeDb()->getContainer("StsCalPar");
    if (!pCalPar)
    {
        Error("HForwardBasicTracker::init()", "Parameter container StsCalPar not found");
        return kFALSE;
    }

    return kTRUE;
}

Bool_t HForwardBasicTracker::reinit()
{
    fLRCutChi2 = pStrawVFPar->getLRCutChi2();
    fTanCut = pStrawVFPar->getTanCut();

    fLRErrU = pStrawVFPar->getLRErrU();

    fCutX = pStrawVFPar->getCutX();
    fCutY = pStrawVFPar->getCutY();

    nMaxBest = pStrawVFPar->getMaxBest();
    fMaxHits = pStrawVFPar->getMaxAllowedHits();

    assert(fMaxHits <= MAX_POSSIBLE_HITS);

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
        {
            HModGeomPar* fmodgeom = pStsGeomPar->getModule(l, m);
            HGeomTransform labTrans = fmodgeom->getLabTransform();
            HGeomCompositeVolume* fMod = fmodgeom->getRefVolume();

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

    fZavg = 0.f;
    HGeomVector g_ref(0, 1, 0);
    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        assert(stsCellsLab[m][0][0] != nullptr);
        fZ0[m] = stsCellsLab[m][0][0]->Z();

        // Loop over layers
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
        {
            HModGeomPar* lvol = pStsGeomPar->getModule(l, m);
            const HGeomTransform& labTrans = lvol->getLabTransform();

            for (Int_t p = 0; p < STS_MAX_PLANES; ++p)
            {
                Int_t plane = HStsCal::getVPlane(m, l, p);

                assert(stsCellsLab[m][l][p] != nullptr);
                fZ[plane] = stsCellsLab[m][l][p]->Z();
                fDz[plane] = fZ[plane] - fZ0[0];
                fCosa[plane] = labTrans.getRotMatrix().getElement(1, 1);
                fSina[plane] = -fabs(labTrans.getRotMatrix().getElement(1, 0));
                fZavg += fZ[plane];

                if (fCosa[plane] < 0)
                {
                    fCosa[plane] = -fCosa[plane];
                    fSina[plane] = -fSina[plane];
                }
            }
        }
    }
    fZavg /= STS_MAX_VPLANES;

    computeMatrix(); // compute system matrices

#ifdef VERBOSE_MODE
    pStrawVFPar->printParams();
#endif

    return kTRUE;
}

void HForwardBasicTracker::clear(Bool_t all)
{
    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        fModuleVectors[m].clear();
    }

    if (all)
    {
        while (!fVectorsFull.empty())
        {
            delete fVectorsFull.front();
            fVectorsFull.pop_front();
        }
    }
}

Int_t HForwardBasicTracker::execute()
{
#ifdef VERBOSE_MODE
    static Int_t event = 0;
    printf("########### EVENT %d\n", event++);
#endif
    // gErrorIgnoreLevel = kInfo; //debug level
    gErrorIgnoreLevel = kWarning; // debug level

    Int_t nhits = pStrawCal->getEntries();

    clear();

    if (nhits > fMaxHits)
    {
        Error("HForwardBasicTracker::execute()", "Event: %d  To many hits for tracking %i > %i",
              gHades->getCurrentEvent()->getHeader()->getEventSeqNumber(), nhits, pStrawVFPar->getMaxAllowedHits());
        return kFALSE;
    }
    getHits();

    makeVectors();

    // Remove vectors with wrong orientation
    checkParams();

    // Match vectors from the first and the second station
    matchVectors();

    // Store tracks
    storeVectors();

    ++n_events;

    return 0;
}

Bool_t HForwardBasicTracker::finalize()
{
    static Int_t ct = 0;

    if (ct == 0)
    {
        printf("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@\n");
        printf("@\n");
        printf("@ Forward Basic Tracker statistics\n");
        printf("@\n");
        printf("@ Analyzed events:             %d\n", n_events);
        printf("@ Events with vectors (OK):    %d  (%.2f%%)\n", n_ok_events, 100. * n_ok_events / n_events);
        printf("@ Total reconstructed vectors: %d\n", n_vecs);
        printf("@ Average vectors/OK event:    %.3f\n", 1. * n_vecs / n_ok_events);
        printf("@\n");
        printf("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@\n");
        ct++;
    }
    return kTRUE;
}

void HForwardBasicTracker::computeMatrix()
{
    // delete existing matrices
    std::map<Int_t, TDecompLU*>::iterator it_lus;
    for (it_lus = fLus.begin(); it_lus != fLus.end(); ++it_lus)
        delete it_lus->second;

    std::map<Int_t, TMatrixDSym*>::iterator it_matr;
    for (it_matr = fMatr.begin(); it_matr != fMatr.end(); ++it_matr)
        delete it_matr->second;

    // Compute system matrices for different hit plane patterns
    Double_t cos2[STS_MAX_VPLANES];
    Double_t sin2[STS_MAX_VPLANES];
    Double_t sincos[STS_MAX_VPLANES];
    Double_t dz2[STS_MAX_VPLANES];

    for (Int_t i = 0; i < STS_MAX_VPLANES; ++i)
    {
        cos2[i] = fCosa[i] * fCosa[i];
        sin2[i] = fSina[i] * fSina[i];
        sincos[i] = fSina[i] * fCosa[i];
        dz2[i] = fDz[i] * fDz[i];
    }

    TMatrixD coef(4, 4);
    Int_t pattMax = 1 << STS_MAX_VPLANES, pattMax1 = 1 << STS_MAX_VPLANES / 2, nDouble = 0, nTot = 0;

    // Loop over doublet patterns
    for (Int_t ipat = 1; ipat < pattMax; ++ipat)
    {
        // Check if the pattern is valid:
        // either all doublets are active or
        // 3 the first ones (for high resolution mode)

        nDouble = 0;
        if (ipat < pattMax1)
        {
            // One station
            for (Int_t j = 0; j < (STS_MAX_VPLANES / 2); j += 2)
            {
                if (ipat & (3 << j))
                    ++nDouble;
                else
                    break;
            }
        }
        else
        {
            // The other station
            Int_t ipat1 = ipat & ((pattMax1 - 1) << (STS_MAX_VPLANES / 2));
            for (Int_t j = 0; j < (STS_MAX_VPLANES / 2); j += 2)
            {
                if (ipat1 & (3 << (j + (STS_MAX_VPLANES / 2))))
                    ++nDouble;
                else
                    break;
            }
        }

        if (nDouble + 0 < (STS_MAX_VPLANES / 2) / 2) continue;

        coef = 0.0;
        for (Int_t i = 0; i < STS_MAX_VPLANES; ++i)
        {
            Bool_t onoff = (ipat & (1 << i));
            if (!onoff) continue;

            Double_t dz = fDz[i];
            Double_t dzz = dz2[i];

            coef(0, 0) += cos2[i];
            coef(0, 1) += sincos[i];
            coef(0, 2) += dz * cos2[i];
            coef(0, 3) += dz * sincos[i];

            coef(1, 0) += sincos[i];
            coef(1, 1) += sin2[i];
            coef(1, 2) += dz * sincos[i];
            coef(1, 3) += dz * sin2[i];

            coef(2, 0) += dz * cos2[i];
            coef(2, 1) += dz * sincos[i];
            coef(2, 2) += dzz * cos2[i];
            coef(2, 3) += dzz * sincos[i];

            coef(3, 0) += dz * sincos[i];
            coef(3, 1) += dz * sin2[i];
            coef(3, 2) += dzz * sincos[i];
            coef(3, 3) += dzz * sin2[i];
        }
        if (coef.Determinant() < 1.e-15) continue;

        ++nTot;

        TDecompLU* lu = new TDecompLU(4);
        lu->SetMatrix(coef);
        lu->SetTol(0.1 * lu->GetTol());
        lu->Decompose();
        fLus.insert(pair<Int_t, TDecompLU*>(ipat, lu));

        TMatrixDSym cov(4);
        cov.SetMatrixArray(coef.GetMatrixArray());
        cov.Invert(); // covar. matrix
        fMatr.insert(pair<Int_t, TMatrixDSym*>(ipat, new TMatrixDSym(cov)));
    }
}

Int_t HForwardBasicTracker::getHits()
{
    // Group hits according to their plane number
    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        for (Int_t l = 0; l < STS_MAX_LAYERS * STS_MAX_PLANES; ++l)
            fHitPl[m][l].clear();
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
            fPair[m][l].clear();
    }

    Int_t nHits = pStrawCal->getEntries();
    Char_t addr_module, addr_layer, addr_udconf;
    Int_t addr_straw, module, layer;

    Int_t fUD[MAX_POSSIBLE_HITS];        // hit ud conf
    Float_t fTime[MAX_POSSIBLE_HITS];    // hit time
    Float_t fTimeMax[MAX_POSSIBLE_HITS]; // hit time

    Int_t accepted_hits = 0;
    for (Int_t i = 0; i < nHits; ++i)
    {
        fStraw[i] = (HStsCal*)pStrawCal->getObject(i);

        if (fStraw[i]->GetUniqueID() != 0) continue; // already used

        ++accepted_hits;

        fStraw[i]->getAddress(addr_module, addr_layer, addr_straw, addr_udconf);
        module = (Int_t)addr_module;
        layer = (Int_t)addr_layer;

        // Plane calculation.
        // Typical double layer has configuration:
        //
        // z ^
        //   |   O O O B       p=1
        //   |  A O O O        p=0
        //
        // where p are planes and O,A,B are straw sumbers. But depending of the plane rotation
        // the straw no. 0 can be either on A or B place. Thus to calculate plane p from a straw
        // number N  p = N % 2 one has to determine which variant it is, using geometry Z
        // coordinate. Larger Z means B is the first otherwise A.

        Float_t z1 = stsCellsLoc[module][layer][0]->Z();
        Float_t z2 = stsCellsLoc[module][layer][1]->Z();

        Int_t plane;
        if (z1 < z2)
            plane = addr_straw % 2; // case A
        else
            plane = 1 - addr_straw % 2; // case B

        Int_t pstraw = addr_straw / 2;

        Int_t vplane = fStraw[i]->getVPlane(0, layer, plane);

#ifdef VERBOSE_MODE_PAIRS
        printf("(%3d)  m=%d l=%d p=%d -> vp=%d  ps=%03d  time=%7.3f  u=%8.3f  z=%8.3f  p=%12f (%8.2f,%8.2f,%8.2f || "
               "%+8.2f,%+8.2f)  "
               " tr=%d\n",
               i, module, layer, plane, vplane, pstraw, fStraw[i]->getTime(), fStraw[i]->getU(), fStraw[i]->getZ(),
               isSimulation ? ((HStsCalSim*)fStraw[i])->getP() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getPx() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getPy() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getPz() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getPx() / ((HStsCalSim*)fStraw[i])->getPz() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getPy() / ((HStsCalSim*)fStraw[i])->getPz() : -100.0,
               isSimulation ? ((HStsCalSim*)fStraw[i])->getTrack() : -1);
#endif

        fHitPl[module][vplane].insert(HitPair(pstraw, i));
        fUD[i] = fStraw[i]->getUpDown();
        fTime[i] = fStraw[i]->getTime();
        fTimeMax[i] = (*pCalPar)[module][layer][addr_straw][addr_udconf].getTimeMax();
    }

    // Merge neighbouring hits from two planes of one layer.
    // If there is no neighbour, takes hit from a single plane
    // (to account for inefficiency)
    Int_t indx1 = 0, indx2 = 0, tube1 = 0, tube2 = 0;

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        // Loop over stations
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
        {
            // Straws in two planes are configured as following:
            // I)
            //
            //      \   | |   / <-- tracks
            // p=1  ( 1)( 2)( 3)...( n)  <-- tube number
            // p=0    ( 1)( 2)( 3)...( n)
            //         \| |/
            //
            // possible pairs for [p=0,p=1] and for 1 < k < n
            // are [k, k] and [k, k+1] -> dk = 0
            //
            // or:
            // II)
            //
            //       \   | |   / <-- tracks
            // p=1    ( 1)( 2)( 3)...( n)  <-- tube number
            // p=0  ( 1)( 2)( 3)...( n)
            //          \| |/
            //
            // possible pairs for [p=0,p=1] and for 1 < k < n
            // are [k, k-1] and [k, k] -> dk = 1
            //
            // Find dk using geometry:
            // Algorithm is as follow:
            // 1. Iterate over p=0
            // 2. check for partners in p=1
            // 3. when all straws from p=0 are done
            // 4. fill remaining straws from p=1

            Float_t x1 = stsCellsLoc[m][l][0]->X();
            Float_t x2 = stsCellsLoc[m][l][1]->X();
            Int_t dk = fabs(x1) < fabs(x2) ? 0 : 1;

            multimap<Int_t, Int_t>::const_iterator it1end = fHitPl[m][l * STS_MAX_PLANES].end(),
                                                   it2end = fHitPl[m][l * STS_MAX_PLANES + 1].end();

            // Loop over doublets
            for (Int_t pass = 0; pass < STS_MAX_PLANES; ++pass)
            {
#ifdef VERBOSE_MODE_PAIRS
                printf("pass %d\n", pass);
#endif
                multimap<Int_t, Int_t>::iterator it1 = fHitPl[m][l * STS_MAX_PLANES].begin(),
                                                 it2 = fHitPl[m][l * STS_MAX_PLANES + 1].begin();

                Bool_t tube1_has_partner = kFALSE; // if tube1 was matched
                Bool_t tube2_has_partner = kFALSE; // if tube2 was matched

                while (it1 != it1end)
                {
                    indx1 = it1->second;
                    tube1 = it1->first;

                    const Int_t ud1 = fUD[indx1];

                    if (it2 == it2end)
                    {
                        if (!tube1_has_partner)
                        {
                            Bool_t badCond = ((pass == 0 and ud1 == 0x2) or (pass == 1 and ud1 != 0x2));
                            if (!badCond) fPair[m][l].push_back(DoubletPair(indx1, -1));
                        }
                        ++it1;
                        tube1_has_partner = kFALSE;

                        continue;
                    }

                    while (it2 != it2end)
                    {
                        indx2 = it2->second;
                        tube2 = it2->first;

                        const Int_t ud2 = fUD[indx2];

                        Bool_t skipThisPair = kFALSE;

                        if (pass == 0)
                        {
                            Bool_t doBreak = kFALSE;
                            Bool_t doNext = kFALSE;

                            // we dont want lower tubes here
                            if (ud1 == 0x2)
                            {
                                ++it1;
                                doBreak = kTRUE;
                            }
                            if (ud2 == 0x2)
                            {
                                ++it2;
                                doNext = kTRUE;
                            }
                            if (doBreak) break;
                            if (doNext) continue;
                        }
                        else if (pass == 1)
                        {
                            Bool_t doBreak = kFALSE;
                            Bool_t doNext = kFALSE;

                            // we dont want upper tubes here
                            if (ud1 == 0x1)
                            {
                                ++it1;
                                doBreak = kTRUE;
                            }
                            if (ud2 == 0x1)
                            {
                                ++it2;
                                doNext = kTRUE;
                            }
                            if (doBreak) break;
                            if (doNext) continue;

                            if (ud1 == 0x3 && ud2 == 0x3)
                            {
                                ++it1;
                                ++it2;
                                break;
                            }
                            if (ud1 == 0x3 && it2 == it2end)
                            {
                                ++it1;
                                break;
                            }
                            if (ud2 == 0x3 && it1 == it1end)
                            {
                                ++it2;
                                continue;
                            }
                        }

                        // difference between tubes
                        Int_t tube_diff = tube1 - tube2 - dk;

#ifdef VERBOSE_MODE_PAIRS
                        if (skipThisPair) printf("  skipped!  ");

                        printf("Testing pair (tube, index, ud): (%d, %d, %d), (%d, %d, %d) dist=%d\n", tube1, indx1,
                               ud1, tube2, indx2, ud2, tube_diff);
#endif
                        if (tube_diff > 0) // tube1 ahead of tube2
                        {
                            if (!tube2_has_partner)
                            {
                                if (!skipThisPair) fPair[m][l].push_back(DoubletPair(-1, indx2));
                            }

                            ++it2; // searching for [k, k] case
                            tube2_has_partner = kFALSE;

                            continue;
                        }

                        if (tube_diff < -1) // tube2 ahead of tube1
                        {
                            if (!tube1_has_partner)
                            {
                                if (!skipThisPair) fPair[m][l].push_back(DoubletPair(indx1, -1));
                            }
                            tube1_has_partner = kFALSE;
                            ++it1; // searching for [k, k] case
                            break;
                        }

#ifdef VERBOSE_MODE_PAIRS
                        printf(" %d %f  %d %f   %f %f   -> %d  %d\n", indx1, fTime[indx1], indx2, fTime[indx2],
                               fTimeMax[indx1], fTimeMax[indx2],
                               fabs(fTime[indx2] - fTime[indx1]) > std::min(fTimeMax[indx1], fTimeMax[indx2]) * 1.2,
                               tube_diff);
#endif
                        if (fabs(fTime[indx2] - fTime[indx1]) > std::min(fTimeMax[indx1], fTimeMax[indx2]) * 1.2)
                            skipThisPair = kTRUE;

                        if (tube_diff == 0) // we have first possible pair
                        {
                            if (!skipThisPair) fPair[m][l].push_back(DoubletPair(indx1, indx2));

                            tube1_has_partner = kTRUE;

                            ++it2; // searching for [k, k] case
                            tube2_has_partner = kFALSE;

                            continue;
                        }

                        if (tube_diff == -1) // we have second possible pair
                        {
                            if (!skipThisPair) fPair[m][l].push_back(DoubletPair(indx1, indx2));

                            tube2_has_partner = kTRUE;

                            ++it1; // searching for [k+1, k] case
                            tube1_has_partner = kFALSE;

                            break;
                        }
                    }
                }

                while (it2 != it2end)
                {
                    indx2 = it2->second;
                    tube2 = it2->first;

                    const Int_t& ud2 = fUD[indx2];

                    if (!tube2_has_partner)
                    {
                        Bool_t badCond = ((pass == 0 and ud2 == 0x2) or (pass == 1 and ud2 != 0x2));
                        if (!badCond) fPair[m][l].push_back(DoubletPair(-1, indx2));
                    }

                    ++it2;
                    tube2_has_partner = kFALSE;
                }
            }
        }
    }

#ifdef VERBOSE_MODE_PAIRS
    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        printf("mod=%2d\n", m);
        for (Int_t l = 0; l < STS_MAX_LAYERS; ++l)
        {
            for (UInt_t d = 0; d < fPair[m][l].size(); ++d)
                printf("d=%d/%ld m=%d 2d=%d   f=%d  s=%d\n", d + 1, fPair[m][l].size(), m, l, fPair[m][l][d].first,
                       fPair[m][l][d].second);
        }
    }
#endif

    return accepted_hits;
}

void HForwardBasicTracker::makeVectors()
{
    // Make vectors for stations
    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        Int_t patt = 0, ndoubl = fPair[m][0].size();

        for (Int_t id = 0; id < ndoubl; ++id)
        {
            Int_t indx1 = fPair[m][0][id].first;
            Int_t indx2 = fPair[m][0][id].second;

            Int_t straw_obj_index[STS_MAX_VPLANES] = {0};
            straw_obj_index[HStsCal::getVPlane(m, 0, 0)] = indx1 >= 0 ? indx1 : -1;
            straw_obj_index[HStsCal::getVPlane(m, 0, 1)] = indx2 >= 0 ? indx2 : -1;
            Bool_t ind1 = indx1 + 1;
            Bool_t ind2 = indx2 + 1;
            patt = ind1;
            patt |= (ind2 << 1);

#ifdef VERBOSE_MODE4
            printf("makeVectors(): m=%d  l=%d  id=%d  patt=%x  idx=%d,%d\n", m, 0, id, patt, indx1, indx2);
#endif
            processDouble(m, 1, patt, straw_obj_index);
        }
    }

#ifdef VERBOSE_MODE4
    printf("makeVectors(): created %ld vectors\n", fModuleVectors[0].size());
#endif
}

void HForwardBasicTracker::processDouble(Int_t m, Int_t l, Int_t patt, Int_t (&straw_obj_index)[STS_MAX_VPLANES])
{
    // Main processing engine (recursively adds doublet hits to the vector)
    Double_t pars[5] = {0.0, 0.0, 0.0, 0.0, fZ0[0]};

    Int_t ndoubl = fPair[m][l].size();

    Int_t pl = HStsCal::getVPlane(m, l, 0);

    for (Int_t id = 0; id < ndoubl; ++id)
    {
        Int_t indx1 = fPair[m][l][id].first;
        Int_t indx2 = fPair[m][l][id].second;

#ifdef VERBOSE_MODE_LRT
        printf("%*cprocessDouble(): m=%d  l=%d  id=%d  patt=%x  pl=%d  idx=%d,%d\n", l, ' ', m, l, id, patt, pl, indx1,
               indx2);
#endif

        HStsCal* hit = (HStsCal*)fStraw[TMath::Max(indx1, indx2)];

        // Check the same views
        Int_t ok = 1;

        for (Int_t i = HStsCal::getVPlane(m, 0, 0); i < pl; ++ ++i)
        {
            Int_t ipl = i;
            if (straw_obj_index[ipl] == -1) continue;
            // if no hit at ipl then check ipl+1
            if (!(patt & (1 << ipl))) ++ipl;

            if (TMath::Abs(fSina[pl] - fSina[ipl]) < 0.01)
            {
                // Check tube difference
                Float_t du12 = abs(hit->getU() - fStraw[straw_obj_index[ipl]]->getU());
                Float_t dz12 = fDz[pl] - fDz[ipl];
                Float_t tan = du12 / dz12;
#ifdef VERBOSE_MODE_LRT
                printf("%*c[%d,%d] -> [(%d,%f,%.1f),(%d,%f,%.1f)] | %f <? %f = TanCut\n", l, ' ', ipl, pl, -1,
                       fStraw[straw_obj_index[ipl]]->getU(), fDz[ipl], hit->getStraw(), hit->getU(), fDz[pl], tan,
                       fTanCut);
#endif
                if (tan > fTanCut)
                {
                    ok = 0;
                    break;
                }
            }
        }

        if (!ok) continue;

        straw_obj_index[HStsCal::getVPlane(m, l, 0)] = indx1 >= 0 ? indx1 : -1;
        straw_obj_index[HStsCal::getVPlane(m, l, 1)] = indx2 >= 0 ? indx2 : -1;

        Bool_t ind1 = indx1 + 1;
        Bool_t ind2 = indx2 + 1;

        Int_t patt_new = patt;
        patt_new |= (ind1 << l * STS_MAX_PLANES); // Set bits
        patt_new |= (ind2 << (l * STS_MAX_PLANES + 1));

        if (l + 1 < (STS_MAX_VPLANES / 2) / STS_MAX_PLANES) { processDouble(m, l + 1, patt_new, straw_obj_index); }
        else
        {
            // Straight line fit of the vector
            Int_t patt1 = patt_new << m * (STS_MAX_VPLANES / 2);
            findLine(patt1, straw_obj_index, pars);
            Double_t chi2 = findChi2(patt1, straw_obj_index, pars, fLRErrU);

#ifdef VERBOSE_MODE4
            printf("%*c+chi2 = %f < %f for patt=%x\n", l, ' ', chi2, fLRCutChi2, patt1);
#endif
            // add vector to the temporary container
            if (chi2 <= fLRCutChi2)
            {
                VectorInfo vi;
                vi.patt = patt1;
                memcpy(&vi.straw_obj_index[0], &straw_obj_index[0], sizeof(vi.straw_obj_index));
                memcpy(&vi.pars[0], pars, sizeof(vi.pars));

                fModuleVectors[m].push_back(vi);
            }
        }
    }
}

/**
 * Fit hits to the straight line
 */
void HForwardBasicTracker::findLine(Int_t patt, const Int_t (&straw_obj_index)[STS_MAX_VPLANES], Double_t* pars)
{
    TVectorD b(4);

    for (Int_t i = 0; i < STS_MAX_VPLANES; ++i)
    {
        Bool_t onoff = patt & (1 << i);
        if (!onoff) continue;

        Int_t pi = straw_obj_index[i];
        b[0] += fStraw[pi]->getU() * fCosa[i];
        b[1] += fStraw[pi]->getU() * fSina[i];
        b[2] += fStraw[pi]->getU() * fCosa[i] * fDz[i];
        b[3] += fStraw[pi]->getU() * fSina[i] * fDz[i];
    }

    // Solve system of linear equations
    Bool_t ok;
    TVectorD solve = fLus[patt]->Solve(b, ok);
    for (Int_t i = 0; i < 4; ++i)
        pars[i] = solve[i];
}

/**
 * Compute chi2 of the fit
 */
Double_t HForwardBasicTracker::findChi2(Int_t patt, const Int_t (&straw_obj_index)[STS_MAX_VPLANES], Double_t* pars,
                                        Float_t error)
{
    Float_t chi2 = 0.;

    for (Int_t i = 0; i < STS_MAX_VPLANES; ++i)
    {
        Bool_t onoff = patt & (1 << i);
        if (!onoff) continue;

        Float_t x = pars[0] + pars[2] * fDz[i];
        Float_t y = pars[1] + pars[3] * fDz[i];
        Float_t u = x * fCosa[i] + y * fSina[i];

        Double_t du = (u - fStraw[straw_obj_index[i]]->getU()) / error;
        chi2 += du * du;
#ifdef VERBOSE_MODE4
        printf("   chi2: p=%02d  chi2=%.3f   | loc=%.2f,%.2f  sina=%.2f  cosa=%.2f  u=%f  Uz=%f  Dz=%.2f  "
               "pars=%.2f,%.2f,%.2f,%.2f  %.2f\n",
               i, du * du, x, y, fSina[i], fCosa[i], u, fStraw[straw_obj_index[i]]->getU(), fDz[i], pars[0], pars[1],
               pars[2], pars[3], pars[4]);
#endif
    }

    return chi2;
}

void HForwardBasicTracker::setTrackId(HForwardCand* vec)
{
    if (!isSimulation) { return; }

    std::map<Int_t, Int_t> ids;
    Int_t nhits = vec->getNofHits();

    // count all associated track ids
    for (Int_t ih = 0; ih < nhits; ++ih)
    {
        HStsCalSim* hit = (HStsCalSim*)fStraw[ih];
        ++ids[hit->getTrack()];
    }

    Int_t currentMax = 0;
    Int_t arg_max = 0;

    std::map<Int_t, Int_t>::iterator it;

    // find the most frequent occurance
    for (it = ids.begin(); it != ids.end(); ++it)
    {
        if (it->second > currentMax)
        {
            arg_max = it->first;
            currentMax = it->second;
        }
    }

    // to have a valid ID more than half of the tracks must have the same ID
    if (currentMax > nhits / 2)
        fillForwardCandSim(vec, arg_max);
    else
        fillForwardCandSim(vec, -1);
}

void HForwardBasicTracker::checkParams()
{
    // Remove vectors with wrong orientation

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
#ifdef VERBOSE_MODE_CHKPARS
        Int_t cnt = fModuleVectors[m].size();
        Int_t nvec = cnt;
        printf("Testing vectors chi2 with m=%d n=%d\n", m, fModuleVectors[m].size());
#endif
        std::vector<VectorInfo>::iterator it = fModuleVectors[m].begin();
        while (it != fModuleVectors[m].end())
        {
#ifdef VERBOSE_MODE_CHKPARS
            printf("[%d] chi2 = %f\n", nvec - cnt, findChi2(it->patt, it->straw_obj_index, it->pars, fLRErrU));
#endif
            const Double_t x = it->pars[0];
            const Double_t y = it->pars[1];
            const Double_t tx = it->pars[2];
            const Double_t ty = it->pars[3];
            const Double_t z = it->pars[4];
            Double_t dTx = tx - x / z;
            Double_t dTy = ty - y / z;
            if (TMath::Abs(dTx) > fCutX || TMath::Abs(dTy) > fCutY)
            {
                it = fModuleVectors[m].erase(it);
#ifdef VERBOSE_MODE_CHKPARS
                --cnt;
#endif
            }
            else { ++it; }
        }
#ifdef VERBOSE_MODE_CHKPARS
        printf(" Vectors after parameter check: m=%d, inp=%d, out=%d\n", m, nvec, cnt);
#endif
    }
}

void HForwardBasicTracker::matchVectors()
{
    Int_t max_n = std::max(fModuleVectors[0].size(), fModuleVectors[1].size());
    Int_t min_n = std::min(fModuleVectors[0].size(), fModuleVectors[1].size());

    if (min_n == 0) return;

    const Int_t N = 2;
    const Int_t kNN = 2;

    Int_t points_n[STS_MAX_MODULES];
    Double_t points_x[STS_MAX_MODULES][max_n];
    Double_t points_y[STS_MAX_MODULES][max_n];
    VectorInfo* points_links[STS_MAX_MODULES][max_n];

    TKDTreeID* kdtree[STS_MAX_MODULES] = {nullptr};

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        Int_t cnt = 0;
        points_n[m] = fModuleVectors[m].size();
        std::vector<VectorInfo>::iterator it = fModuleVectors[m].begin();
        while (it != fModuleVectors[m].end())
        {
            TVector2 pp = HForwardTools::getPlaneIntersect(HGeomVector(it->pars[0], it->pars[1], it->pars[4]),
                                                           HGeomVector(it->pars[2], it->pars[3], 1.), fZavg);
            points_x[m][cnt] = pp.X();
            points_y[m][cnt] = pp.Y();
            points_links[m][cnt] = &(*it);
            ++cnt;
            ++it;
        }
        kdtree[m] = new TKDTreeID(points_n[m], N, 1);
        kdtree[m]->SetData(0, points_x[m]);
        kdtree[m]->SetData(1, points_y[m]);
        kdtree[m]->Build();
    }

    // see
    // https://stackoverflow.com/questions/5077318/given-two-large-sets-of-points-how-can-i-efficiently-find-pairs-that-are-near
    for (Int_t i = 0; i < points_n[0]; ++i)
    {
        Double_t p1[N] = {points_x[0][i], points_y[0][i]};
        Int_t ind1[kNN], ind2[kNN];
        Double_t dist1[kNN], dist2[kNN];

        kdtree[1]->FindNearestNeighbors(p1, kNN, ind2, dist2);
        for (Int_t j = 0; j < kNN; ++j)
        {
            if (ind2[j] == -1) break;

            Double_t p2[N] = {points_x[1][ind2[j]], points_y[1][ind2[j]]};
            kdtree[0]->FindNearestNeighbors(p2, kNN, ind1, dist1);
            for (Int_t k = 0; k < kNN; ++k)
            {
                if (ind1[k] == -1) break;

                if (i == ind1[k])
                {
                    Int_t patt = points_links[0][ind1[k]]->patt | points_links[1][ind2[j]]->patt;

                    Int_t straw_obj_index[STS_MAX_VPLANES];
                    memcpy(straw_obj_index, points_links[0][ind1[k]]->straw_obj_index, sizeof(straw_obj_index) / 2);
                    memcpy(straw_obj_index + STS_MAX_VPLANES / 2,
                           points_links[1][ind2[j]]->straw_obj_index + STS_MAX_VPLANES / 2,
                           sizeof(straw_obj_index) / 2);

                    Double_t pars[5] = {0.0, 0.0, 0.0, 0.0, fZ0[0]};
                    findLine(patt, straw_obj_index, pars);

                    Double_t chi2 = findChi2(patt, straw_obj_index, pars, fLRErrU);
#ifdef VERBOSE_MODE_LRT
                    printf("Chi2 vs cut: %f  %f\n", chi2, fLRCutChi2);
#endif
                    if (chi2 < fLRCutChi2 * 1.5) addVector(patt, straw_obj_index, chi2, pars);
                }
            }
        }
    }

    for (Int_t m = 0; m < STS_MAX_MODULES; ++m)
    {
        delete kdtree[m];
    }
}

HForwardCand* HForwardBasicTracker::addVector(Int_t patt, const Int_t (&straw_obj_index)[STS_MAX_VPLANES],
                                              Double_t chi2, Double_t* pars)
{
// Add vector to the temporary container (as HForwardCand)
#ifdef VERBOSE_MODE4
    printf("      + addVector():  patt=%x  chi2=%f\n", patt, chi2);
    printf("        pars=%f %f %f %f  %f\n", pars[0], pars[1], pars[2], pars[3], pars[4]);
#endif

    HForwardCand* track;
    if (isSimulation)
        track = (HForwardCand*)new HForwardCandSim();
    else
        track = new HForwardCand();

    track->setChi2(chi2);
    HForwardTools::setHadesParams(track, pars);
    TMatrixDSym cov(*fMatr[patt]);
    cov *= (fLRErrU * fLRErrU);
    track->setCovMatrix(cov);

    Float_t tot = 0.0;
    Int_t idx = 0;
    for (Int_t ipl = 0; ipl < STS_MAX_VPLANES; ++ipl)
    {
        if (!(patt & (1 << ipl))) continue;

        Int_t pi = straw_obj_index[ipl];
        track->addStsHit(pi);

        // printf("%x  %d\n", hit, pi);
        tot += fStraw[pi]->getTot();
        track->setStsDriftRadius(idx, 2.5);

        if (isSimulation)
        {
            HStsCalSim* hit = (HStsCalSim*)fStraw[pi];
            ((HForwardCandSim*)track)->addGeantHitTrack(hit->getTrack());
        }
        idx++;
    }
    track->setTot(tot);

    Int_t ndf = (track->getNofHits() > 5) ? track->getNofHits() - 4 : 1;
    track->setNDF(ndf);

    if (isSimulation) { ((HForwardCandSim*)track)->calcGeantCorrTrackIds(); }

    fVectorsFull.push_back(track);

    return track;
}

void HForwardBasicTracker::storeVectors()
{
    // Store vectors into category
    HLocation loc;
    loc.set(1, 0);

#ifdef VERBOSE_MODE
    Int_t vec_cnt = 0;
#endif

    Bool_t has_vector = kFALSE;
    std::list<HForwardCand*>::iterator it = fVectorsFull.begin();
    while (it != fVectorsFull.end())
    {
        HForwardCand* cand = (HForwardCand*)pForwardCand->getNewSlot(loc);

        if (cand)
        {
            if (isSimulation)
                cand = (HForwardCand*)new (cand) HForwardCandSim(*(HForwardCandSim*)(*it));
            else
                cand = new (cand) HForwardCand(*(*it));

            cand->SetUniqueID(std::distance(fVectorsFull.begin(), it));
            setTrackId(cand);

            has_vector = kTRUE;
            ++n_vecs;

#ifdef VERBOSE_MODE
            printf("HForwardBasicTracker::storeVectors\n");
            if (isSimulation)
                ((HForwardCandSim*)cand)->print();
            else
                cand->print();
            ++vec_cnt;
#endif
        }
        ++it;
    }
    if (has_vector) ++n_ok_events;
#ifdef VERBOSE_MODE
    printf("Found %d vectors from %d hits.\n", vec_cnt, pStrawCal->getEntries());
#endif
}

Bool_t HForwardBasicTracker::fillForwardCandSim(HForwardCand* cand, Int_t track)
{
    if (!isSimulation) return kFALSE;

    HForwardCandSim* part = static_cast<HForwardCandSim*>(cand);
    if (part && pKine) { return HForwardTools::fillVirtualCandSimInfo(part, pKine, track); }
    return kTRUE;
}