//*-- 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 "hforwardfinetracker.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 <Math/Factory.h>
#include <Math/Functor.h>
#include <Math/Minimizer.h>

#include <cassert>
#include <iterator>

// #define VERBOSE_MODE
// #define VERBOSE_MODE_FRPC
// #define VERBOSE_MODE_FILTER
// #define VERBOSE_FIT

HGeomVector v_straws_g[STS_MAX_VPLANES];         // straw axis vector? FIXME
const HGeomVector* v_straws_Pi[STS_MAX_VPLANES]; // straw center position? FIXME not used?
Float_t v_straws_r[STS_MAX_VPLANES];             // straw radius
Float_t v_sigma2 = 0.0;
Float_t v_Z = 0.0;
UInt_t v_patt = 0x0;

Double_t fcn(const Double_t* par)
{
    HGeomVector track_P(par[0], par[1], par[4]);
    HGeomVector track_g(par[2], par[3], 1);

#ifdef VERBOSE_FIT
    track_P.print();
    track_g.print();
#endif
    Double_t f = 0.0;
    Int_t n = 0;
    for (Int_t i = 0; i < STS_MAX_VPLANES; ++i)
    {
        Bool_t onoff = v_patt & (1 << i);
        if (!onoff) continue;

        HGeomVector g_cross = track_g.vectorProduct(v_straws_g[i]);
        HGeomVector P_dist = track_P - *v_straws_Pi[i];

        Double_t DCA = fabs(g_cross.scalarProduct(P_dist) / g_cross.length());
        Double_t dr = DCA - v_straws_r[n];

        f += (dr * dr / v_sigma2);
#ifdef VERBOSE_FIT
        printf("[%2d] v_straws_r = %f  DCA=%f  dr=%f   _f=%f\n", i, v_straws_r[n], DCA, dr, (dr * dr / v_sigma2));
#endif
        ++n;
    }
    if (n > 0)
        f /= n;
    else
        f = 99999999999.9;

    return f;
}

Float_t HForwardFineTracker::minimize(Double_t* pars)
{
    // Set the free variables to be minimized!
    minimizer->SetVariable(0, "x", pars[0], step[0]);
    minimizer->SetVariable(1, "y", pars[1], step[1]);
    minimizer->SetVariable(2, "Ty", pars[2], step[2]);
    minimizer->SetVariable(3, "Ty", pars[3], step[3]);
    minimizer->SetVariable(4, "z", pars[4], step[4]);

    minimizer->Minimize();

    const Double_t* xs = minimizer->X();
    Double_t min_chi2 = fcn(xs);
#ifdef VERBOSE_MODE
    cout << "Minimum: f(" << xs[0] << "," << xs[1] << "," << xs[2] << "," << xs[3] << "): " << min_chi2 << endl;
#endif
    for (size_t i = 0; i < 5; ++i)
        pars[i] = xs[i];

    return min_chi2;
}

HForwardFineTracker::HForwardFineTracker(const Text_t* name, const Text_t* title)
    : HReconstructor(name, title), disable_tof_refit(kFALSE), pStrawCal(nullptr), pFRpcHits(nullptr),
      pForwardCand(nullptr), pStrawVFPar(nullptr), isSimulation(kFALSE), n_events(0), n_ok_events(0), n_vecs(0)

{
    minimizer = ROOT::Math::Factory::CreateMinimizer("Minuit", "Migrad");

    minimizer->SetMaxFunctionCalls(100);
    minimizer->SetMaxIterations(500);
    minimizer->SetTolerance(0.01);

    minimizer_functor = new ROOT::Math::Functor(&fcn, 5);
    for (Int_t i = 0; i < 5; i++)
        step[i] = 0.0001;

    minimizer->SetFunction(*minimizer_functor);
}

HForwardFineTracker::~HForwardFineTracker()
{
    delete minimizer;
    delete minimizer_functor;

    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];
            }
}

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

    HFRpcDetector* pFRpc = (HFRpcDetector*)(gHades->getSetup()->getDetector("FRpc"));
    if (!pFRpc)
    {
        Error("HForwardFineTracker::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("HForwardFineTracker::init()", "HStsCal input missing");
        return kFALSE;
    }

    // straw hits
    pFRpcHits = gHades->getCurrentEvent()->getCategory(catFRpcHit);
    if (!pFRpcHits)
    {
        Error("HForwardFineTracker::init()", "HFRpcHit input missing");
        //         return kFALSE;
    }

    pForwardCand = gHades->getCurrentEvent()->getCategory(catForwardCand);
    if (!pForwardCand)
    {
        Error("HForwardFineTracker::init()", "ForwardCand input missing");
        return kFALSE;
    }

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

    // create the parameter container
    pStsGeomPar = (HStsGeomPar*)gHades->getRuntimeDb()->getContainer("StsGeomPar");
    if (!pStsGeomPar)
    {
        Error("HForwardFineTracker::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;
            }

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

    return kTRUE;
}

Bool_t HForwardFineTracker::reinit()
{
    fHRCutChi2 = pStrawVFPar->getHRCutChi2();

    fHRErrU = pStrawVFPar->getHRErrU();

    fFRpcWindowU = pStrawVFPar->getFRpcTimeWindowU();
    fFRpcWindowL = pStrawVFPar->getFRpcTimeWindowL();

    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;
                HGeomVector* p = stsCellsLab[m][l][c];
                *p = fVol->getTransform().getTransVector();
                *p = labTrans.transFrom(*p);
            }
        }
    }

    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);

                // Calculate axis vector direction for straws in a plane.
                // All straws have in a plane (layer) have the same vector direction.
                v_straws_g[plane] = labTrans.getRotMatrix() * g_ref;
            }
        }
    }

    v_sigma2 = fHRErrU * fHRErrU;

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

    return kTRUE;
}

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

    // get Event vertex
    HVertex vertex = gHades->getCurrentEvent()->getHeader()->getVertexReco();
    if (vertex.getChi2() < 0)
    { // vertex reco has not been reconstructed: fall back solution : Vertex cluster
        vertex = gHades->getCurrentEvent()->getHeader()->getVertexCluster();
        with_vertex_cluster = kTRUE;
    }
    else { with_vertex_cluster = kFALSE; }

    primary_vertex = vertex.getPos();

    // filter best tracks
    filterTracks();

    // Select final tracks - remove ghosts
    selectTracks();

    // Store tracks
    storeVectors();

    ++n_events;

    return 0;
}

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

    if (ct == 0)
    {
        printf("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@\n");
        printf("@\n");
        printf("@ Forward Fine 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 HForwardFineTracker::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*)pStrawCal->getObject(vec->getStsHitIndex(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);
}

class HForwardChi2Filter : public HFilter
{
    Bool_t check(TObject* obj)
    {
        HForwardCand* tr = dynamic_cast<HForwardCand*>(obj);
        if (tr->getChi2() == 0.0) return kFALSE;
        return kTRUE;
    }
};

/**
 * Refit in HR mode and filter tracks which fail HR chi2 test.
 */
void HForwardFineTracker::filterTracks()
{
    // Remove ghost tracks (having at least N the same hits (i.e. fired tubes))
    Int_t n_tracks = pForwardCand->getEntries();

    for (Int_t i = 0; i < n_tracks; ++i)
    {
        HForwardCand* track = dynamic_cast<HForwardCand*>(pForwardCand->getObject(i));

        if (!matchFRpcHit(track))
        {
            if (isSimulation) ((HForwardCandSim*)track)->setGeantTrackRpc(-1);
        }

        // Refit whole line
        Double_t pars[5] = {0.0, 0.0, 0.0, 0.0, fZ0[0]};
        TMatrixDSym cov(4);
        Double_t chi2 = refit(track, pars, &cov);
        Info("Select tracks", "Track chi2: %f", chi2);

#ifdef VERBOSE_MODE_FILTER
        printf("###  FUNC: %s  LINE:  %d\n", __func__, __LINE__);
        printf("New pars: %f  %f    %f  %f   %f\n", pars[0], pars[1], pars[2], pars[3], pars[4]);
        printf("Checking chi2 %f (/ndf %f) for limit %f\n", chi2, chi2 / track->getNDF(), fHRCutChi2);
        track->print();
#endif
        if (chi2 / track->getNDF() > fHRCutChi2)
        {
            track->setChi2(0.0);
#ifdef VERBOSE_MODE_FILTER
            printf(" Track discarded because chi2 %f > %f\n", chi2 / track->getNDF(), fHRCutChi2);
#endif
            continue;
        }
        else { track->setChi2(chi2); }

        cov *= (fHRErrU * fHRErrU);
        HForwardTools::setParams(track, pars);
        track->setCovMatrix(cov);

        if (isSimulation)
        {
            Int_t hidx = 0;
            HStsCalSim* hit = nullptr;
            Float_t px, py, pz, x, y, z;

            hidx = track->getStsHitIndex(0);
            hit = (HStsCalSim*)pStrawCal->getObject(hidx);
            hit->getP(px, py, pz);
            hit->getHitPos(x, y, z);
            if (isSimulation)
            {
                HForwardCandSim* trs = (HForwardCandSim*)track;
                trs->setGeantPx1(px);
                trs->setGeantPy1(py);
                trs->setGeantPz1(pz);
                trs->setGeantX1(x);
                trs->setGeantY1(y);
                trs->setGeantZ1(z);
            }

            hidx = track->getStsHitIndex(track->getNofHits() - 1);
            hit = (HStsCalSim*)pStrawCal->getObject(hidx);
            hit->getP(px, py, pz);
            hit->getHitPos(x, y, z);
            if (isSimulation)
            {
                HForwardCandSim* trs = (HForwardCandSim*)track;
                trs->setGeantPx2(px);
                trs->setGeantPy2(py);
                trs->setGeantPz2(pz);
                trs->setGeantX2(x);
                trs->setGeantY2(y);
                trs->setGeantZ2(z);
            }
        }
    }

    HForwardChi2Filter bad_chi2_filter;
    pForwardCand->filter(bad_chi2_filter);

    n_tracks = pForwardCand->getEntries();
    if (n_tracks)
    {
        n_ok_events++;
        n_vecs += n_tracks;
    }
}

void HForwardFineTracker::storeVectors()
{
    Int_t n_tracks = pForwardCand->getEntries();

    for (Int_t i = 0; i < n_tracks; ++i)
    {
        setTrackId((HForwardCand*)pForwardCand->getObject(i));
    }
}

class HForwardSameStrawsFilter : public HFilter
{
public:
    HForwardSameStrawsFilter(HForwardCand* reference, HCategory* strawcal) : ref(reference), pStrawCal(strawcal)
    {
        Int_t nhits1 = ref->getNofHits();
        std::fill_n(planes, STS_MAX_VPLANES, -1);

        for (Int_t iv = 0; iv < nhits1; ++iv)
        {
            Int_t idx1 = ref->getStsHitIndex(iv);
            HStsCal* hit = (HStsCal*)pStrawCal->getObject(idx1);
            Int_t ipl = hit->getVPlane();
            planes[ipl] = idx1;
        }
    }

    Bool_t check(TObject* obj)
    {
        HForwardCand* tr = dynamic_cast<HForwardCand*>(obj);

        if (tr == ref) return kTRUE;

        if (tr->getChi2() == 0.0) return kTRUE;

        // We can filter out only tracks which are worse than reference
        if (tr->getChi2() / tr->getNDF() < ref->getChi2() / ref->getNDF()) return kTRUE;

        Int_t nhits2 = tr->getNofHits(), nover = 0;

        for (Int_t iv = 0; iv < nhits2; ++iv)
        {
            Int_t idx2 = tr->getStsHitIndex(iv);
            HStsCal* hit = (HStsCal*)pStrawCal->getObject(idx2);
            Int_t ipl = hit->getVPlane();
            if (planes[ipl] < 0) continue;
            if (planes[ipl] == idx2)
            {
                ++nover;
                if (nover >= nMax)
                {
                    Info("Select tracks", "Track quality: qual1 %f, qual2 %f, overlaps: %i",
                         ref->getChi2() / ref->getNDF(), tr->getChi2() / tr->getNDF(), nover);

                    tr->setChi2(0.0);
                    return kTRUE;
                }
            }
        }

        return kTRUE;
    }

    HForwardCand* ref;
    HCategory* pStrawCal;
    Int_t planes[STS_MAX_VPLANES];
    static const Int_t nMax = 5;
};

struct Chi2NDSSorter
{
    Bool_t operator()(const HForwardCand* c1, const HForwardCand* c2)
    {
        return c1->getChi2() / c1->getNDF() < c2->getChi2() / c2->getNDF();
    }
};

/**
 * Finds similar tracks and rejects those which share multiple hits and have worse Chi2.
 */
void HForwardFineTracker::selectTracks()
{
    // Remove ghost tracks (having at least N the same hits (i.e. fired tubes))
    Int_t n_tracks = pForwardCand->getEntries();

    std::vector<HForwardCand*> cands(n_tracks);
    for (Int_t i = 0; i < n_tracks; ++i)
    {
        cands.push_back(dynamic_cast<HForwardCand*>(pForwardCand->getObject(i)));
    }

    for (Int_t i = 0; i < n_tracks; ++i)
    {
        HForwardCand* track = dynamic_cast<HForwardCand*>(pForwardCand->getObject(i));
        if (track->getChi2() == 0.0) continue;

        HForwardSameStrawsFilter same_straws(track, pStrawCal);
        pForwardCand->filter(same_straws);
    }

    HForwardChi2Filter bad_chi2_filter;
    pForwardCand->filter(bad_chi2_filter);
}

Double_t HForwardFineTracker::refit(HForwardCand* track, Double_t* pars, TMatrixDSym* cov)
{
    // Fit of hits to the straight line

    Short_t tofrec = track->getTofRec();
    Double_t tof = track->getTof();

    HStsCal* hits[STS_MAX_VPLANES] = {nullptr};
    Int_t nhits = track->getNofHits();
    UInt_t patt = 0;

    for (Int_t i = 0; i < nhits; ++i)
    {
        HStsCal* hit = (HStsCal*)pStrawCal->getObject(track->getStsHitIndex(i));
        Int_t ipl = hit->getVPlane();
        hits[ipl] = hit;
        patt |= (1 << ipl);
    }

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

        HStsCal* hit = hits[i];

        if (tofrec != 0)
        {
            HStsCalParCellKind& p_cell =
                (*pCalPar)[hit->getModule()][hit->getLayer()][hit->getStraw()][hit->getUpDown()];
            Float_t dr_pars[11];
            p_cell.getData(dr_pars);

            Float_t dr = hit->calcRadius(tof, dr_pars);
            if (dr < 0.0) dr = 0.0;
            track->setStsDriftRadius(idx, dr);
        }

        v_straws_r[i] = track->getStsDriftRadius(idx);
        v_straws_Pi[i] = stsCellsLab[hit->getModule()][hit->getLayer()][hit->getStraw()];
    }

    if (disable_tof_refit) track->setTofRec(2);

    HForwardTools::getParams(track, pars);
#ifdef VERBOSE_MODE
    printf("*********** Refiting ***********\n");
    track->print();
    printf("  PARS: %f %f %f %f %f\n", pars[0], pars[1], pars[2], pars[3], pars[4]);
#endif
    return minimize(pars);
}

HFRpcHit* HForwardFineTracker::matchFRpcHit(HForwardCand* track)
{
    // Match RPC hit to the track in straws
    if (!pFRpcHits) return nullptr;

    Int_t nHits = pFRpcHits->getEntries();

    HFRpcHit* sel_rpc_hit = nullptr;
    Int_t frpcidx = -1;

    if (nHits > 0)
    {
        Double_t min_dist = pStrawVFPar->getFRpcMatchDist();

        Double_t pars[5];
        HForwardTools::getParams(track, pars);
        for (Int_t i = 0; i < nHits; ++i)
        {
            HFRpcHit* hit = (HFRpcHit*)pFRpcHits->getObject(i);

            Double_t dist = HForwardTools::getFRpcDistance(track, hit);

#ifdef VERBOSE_MODE_FRPC
            printf(" (%02d/%02d) RPCHIT   x=%f y=%f z=%.3f tof=%f    for %f,%f (%f,%f) -> dist=%f < %f\n", i, nHits,
                   hit->getX(), hit->getY(), hit->getZ(), hit->getTof(), pars[0], pars[1], pars[2], pars[3], dist,
                   min_dist);
#endif

            if (dist < min_dist)
            {
                // Int_t sel_hit = i;
                sel_rpc_hit = hit;
                frpcidx = i;
                min_dist = dist;
            }
        }
    }

    if (sel_rpc_hit)
    {
        track->setTofRec(1);
        track->setFRpcHitIndex(frpcidx);
        track->setTof(sel_rpc_hit->getTof());
        HForwardTools::correctPathLength(track, primary_vertex, sel_rpc_hit);
        /* Mark that track uses cluster vertex not reco vertex */
        if (with_vertex_cluster) track->setFlagBit(HForward::kVertexCluster);
        HForwardTools::correctPathTof(track, 0.0, sel_rpc_hit, fFRpcWindowL, fFRpcWindowU);
        track->calc4vectorProperties(HPhysicsConstants::mass(14));
        // if (isSimulation) ((HForwardCandSim *)track)->setGeantTrackRpc(sel_rpc_hit->getTrack()); FIXME
    }
    else
    {
        track->setTofRec(0);
        track->setTof(-1.0f);
        track->setDistance(-1.0f);
        track->setFlagBit(HForward::kErrorTofRec);
        track->setFlagBit(HForward::kIsRejected);
    }
    return sel_rpc_hit;
}

Bool_t HForwardFineTracker::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;
}