#pragma once

//
// Adapted from search/relevance/ann_features.h
//

#include <saas/rtyserver/search/cgi/rty_ranking_opts.h>

#include <saas/rtyserver/factors/feature.h>

#include <search/relevance/ann_features.h>

#include <kernel/groupattrs/metainfo.h>

using TRegionWithPrio = std::pair<TCateg, ui32>;

struct TRTYAllRegions {
    TVector<TRegionWithPrio> Filter;

public:
    TRTYAllRegions();
    TRTYAllRegions(TCateg relevGeo, TCateg relevCountry, const TVector<TCateg>& allRegions);

public:
    static TVector<TCateg> MakeAllRegions(const NGroupingAttrs::TMetainfo* geoa, TCateg relevGeo);
};

class TRTYQuorumAnnFeaturesCalcer: public IPosGator {
public:
    TRTYQuorumAnnFeaturesCalcer(const IAnnDataReader*, const TVector<float>& wordWeights, const TVector<TCateg>& allRegions, const TRTYRankingOpts& rankingOpts);
    ~TRTYQuorumAnnFeaturesCalcer();

    void InitNextDoc(TNewDocParams* params) override;
    void AddPositions(TFullPositionEx* pos, size_t count, ERelevanceType rt) override;
    void CalcFeatures(TPosGatorCalcFeaturesParams& params) override;

private:
    void AddPosition(const TFullPositionEx*);

private:
    THolder<NIndexAnn::IDocDataIterator> AnnIter;
    bool HasDoc = false;
    TCateg AnnRelevRegion = TCateg{};
    TRTYAllRegions RegionFilter;

    const IAnnDataReader* Params;
    ui32 DocId;
    TSentenceOffsets Offsets;

    static const ui32 InvalidBreak = -1;

    THolder<class TBclmTracker> NHopBclmTracker;

    TAnnUTrackerWrapper<NIndexAnn::DT_CORRECTED_CTR, 1> CorrCtrTracker;
};

// a helper class for factorann
class TRTYRegStreamValues {
    TBitMap<NIndexAnn::DT_NUM_TYPES> Inited;
    std::array<ui32, NIndexAnn::DT_NUM_TYPES> Values;
    std::array<ui32, NIndexAnn::DT_NUM_TYPES> Prios;
public:
    bool HasValue(NIndexAnn::EDataType dataType) const;
    ui32 GetValue(NIndexAnn::EDataType dataType) const;
    size_t Size() const;
    bool Empty() const;
    void FillFrom(NIndexAnn::IDocDataIterator& iter, ui32 doc, ui32 breuk, const TVector<TRegionWithPrio>& regionChain);

    const TBitMap<NIndexAnn::DT_NUM_TYPES>& GetFlags() const {
        return Inited;
    }
};

class TRTYFactorAnnFeaturesCalcer: public NRTYFeatures::ISlicedCalcer {
private:
    class TImpl;
    THolder<class TImpl> Impl;
    TFactorView ViewMain;

public:
    static const size_t NUM_HIT_POSITIONS = TFactorAnnFeaturesCalcer::NUM_HIT_POSITIONS;

    TRTYFactorAnnFeaturesCalcer();
    ~TRTYFactorAnnFeaturesCalcer();

    void ResetRequest(const NIndexAnn::TReaders& readers,
                      const TRichRequestNode& root,
                      const TVector<TCateg>& allRegions,
                      const TRTYRankingOpts& rankingOpts);

    void Bind(TFactorStorage& storage) override;

    void CalcFeatures(ui32 docId) override;
};
