#include <saas/searchproxy/common/nameextractors.h>
#include <saas/searchproxy/common/cgi.h>

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

#include <search/meta/context.h>

#include "rank.h"

namespace {
    struct TCopyMetaRelevance {
        bool operator()(TMergedDoc& doc) {
            doc.SetMetaRelevance(doc->GetRelevance());
            return true;
        }
    };
} // namespace

namespace NRTYSearchProxy {
    class TDefaultMetaRanker: public IMetaRanker {
    public:
        void CalcMetaRelevance(TMetaGroupingId& gId, const std::pair<size_t, size_t>* /*clients*/) override {
            TMetaGrouping* grouping = gId.Grouping;
            TCopyMetaRelevance f;
            ForEachDoc(grouping, f);
            return;
        }

        void DumpProps(TSearcherProps* /*props*/) override {
        }

        const TRankModelDescr* GetRankModelDescr() const override {
            return nullptr;
        }

        void AdjustFactorsParams(TSourceInfo* /*si*/) override {
        }

        void AdjustFactorsParamsFinally(TMetaSearchContext& /*context*/) override {
        }

        void AdjustFactorsRequests(const IMetaRearrangeContext::TAdjustParams& /*ap*/) override {
        }

        void AdjustFactorsAfterSearch(TMetaGroupingId& /*gId*/) override {
        }

        TString GetReqWebFullFastRankModelName(const size_t* /*client*/) const override {
            return {};
        }

        TString GetReqWebHastRankModelName(const size_t* /*client*/) const override {
            return {};
        }

        void ComputeDumpModels(TMetaGrouping* /*grouping*/, const TVector<TString>& /*dumpModels*/) const override {
        }
    };

    class TFactorsConfigMetaRanker: public TDefaultMetaRanker {
    public:
        TFactorsConfigMetaRanker(TMetaSearchContext* context, const NRTYFactors::TRankModelHolder* rankModel)
            : Ctx(context)
            , RankModel(rankModel) { }

        void CalcMetaRelevance(TMetaGroupingId& gId, const std::pair<size_t, size_t>* /*clients*/) override {
            TMetaGrouping* grouping = gId.Grouping;

            if (Ctx->RP().SP.SortMethod != SORT_BY_RELEV) {
                TCopyMetaRelevance f;
                ForEachDoc(grouping, f);
                return;
            }

            TVector<TMergedDoc*> docs(Reserve(grouping->Size()));
            TVector<float*> docFactors(Reserve(grouping->Size()));
            for (size_t i = 0; i < grouping->Size(); ++i) {
                TMetaGroup& g = grouping->GetMetaGroup(i);
                for (size_t j = 0; j < g.ItemCount(); ++j) {
                    TMergedDoc& doc = g.MetaDocs[j];
                    TFactorStorage* fs = doc.RankingFactors();
                    TFactorView factorView = fs->CreateView();
                    docs.push_back(&doc);
                    docFactors.push_back(factorView.GetRawFactors());
                }
            }
            if (docs.empty()) {
                return;
            }

            // TODO: calc meta features here

            TVector<float> resultMetaRelevance(docs.size(), 0.0f);
            TMaybe<size_t> matrixNetIdx = RankModel->GetMetaMatrixNetIndex();
            TMaybe<size_t> polynomIdx = RankModel->GetMetaPolynomIndex();
            RankModel->MultiCalc(docFactors.data(), resultMetaRelevance.data(), resultMetaRelevance.size(), nullptr, matrixNetIdx, polynomIdx);
            for (size_t i = 0; i < docs.size(); ++i) {
                TMergedDoc& doc = *docs[i];
                doc.SetTrueRelevance(resultMetaRelevance[i]);
                doc.SetMetaRelevance(100000000 + doc.GetTrueRelevance() * 10000000);

                if (auto* attrs = doc->FirstStageAttrValues()) {
                    if (matrixNetIdx.Defined()) {
                        attrs->GetOrAddValues("MetaMatrixNet").AssignSingleValue(ToString(docFactors.data()[i][*matrixNetIdx]));
                    }
                    if (polynomIdx.Defined()) {
                        attrs->GetOrAddValues("MetaPolynom").AssignSingleValue(ToString(docFactors.data()[i][*polynomIdx]));
                    }
                }
            }
        }

        void DumpProps(TSearcherProps* props) override {
            if(Ctx->MetaRankingOptions().Enabled() && Ctx->RP().SP.SortMethod == SORT_BY_RELEV && RankModel) {
                props->Set("MetaModelName", RankModel->GetName());
                if (RankModel->GetRankModel()->HasPolynom()) {
                    props->Set("MetaModelFormula", Encode(*RankModel->GetRankModel()->Polynom()->Descr));
                }
            }
        }

    private:
        TMetaSearchContext* Ctx;
        const NRTYFactors::TRankModelHolder* RankModel;
    };

    class TMetaRankEnv : public IMetaRankEnv {
    public:
        TMetaRankEnv(const TServiceConfig& serviceConfig)
            : ServiceConfig(serviceConfig) { }

        void InitMetaFeatures(const TSearchConfig& /*config*/, TAtomicSharedPtr<NRearr::TDataRegistry> /*dataRegistry*/) override {
        }

        TAutoPtr<IMetaRanker> CreateRanker(TMetaSearchContext* ctx) const override {
            TAtomicSharedPtr<NRTYFactors::TConfig> factors = ServiceConfig.GetFactorsConfig();
            if (!factors) {
                DEBUG_LOG << "No factors config for the service: " << TServiceNameExtractor::ExtractServiceName(ctx->CgiParam) << Endl;
                return new TDefaultMetaRanker();
            }

            const NRTYFactors::TRankModelHolder* rankModel = factors->GetL3RankModel(NSearchProxyCgi::GetL3FormulaName(ctx->CgiParam));
            if (rankModel == nullptr) {
                return new TDefaultMetaRanker();
            }
            return new TFactorsConfigMetaRanker(ctx, rankModel);
        }

        const TRankModels* GetRankModels() const override {
            return nullptr;
        }

        const TRankModelsMapFactory* GetRankModelsMapFactory() const override {
            return nullptr;
        }

        const TMetaFeatures* GetMetaFeatures() const override {
            return nullptr;
        }

        TString GetModelsMD5() const override {
            return nullptr;
        }

        TString TryGetModelsMD5(IThreadPool& /*queue*/) const override {
            return {};
        }

        TString GetModelsSvnVersion() const override {
            return {};
        }

        TString GetModelsSandboxTaskId() const override {
            return {};
        }

        TString GetModelsSandboxResourceId() const override {
            return {};
        }

        TString GetModelsBuildTime() const override {
            return {};
        }

        const IModelsArchiveReader* GetArchiveReader() const override {
            return nullptr;
        }

    private:
        const TServiceConfig& ServiceConfig;
    };

    TAutoPtr<IMetaRankEnv> CreateMetaRankEnv(const TServiceConfig& config) {
        return new TMetaRankEnv(config);
    }
} // namespace NRTYSearchProxy
