#include "dssm_factors.h"
#include "dssm_model_ops.h"
#include <saas/rtyserver/search/features/protos/dssm_models.pb.h>

#include <saas/rtyserver/search/cgi/rty_external_cgi.h>
#include <saas/rtyserver/search/external_search/rty_index_text_data.h> // we include external_search without peerdiring - same as saas/rtyserver/search/cgi - let's call it a trick

#include <kernel/dssm_applier/nn_applier/lib/layers.h>
#include <kernel/factor_storage/factor_storage.h>

#include <search/meta/data_registry/data_registry.h>

#include <library/cpp/protobuf/util/pb_io.h>
#include <library/cpp/logger/global/global.h>
#include <library/cpp/string_utils/url/url.h>

#include <util/folder/path.h>
#include <util/generic/yexception.h>
#include <util/stream/output.h>
#include <util/string/subst.h>
#include <util/string/join.h>


namespace NRTYFeatures {
    //
    // Static metadata initialization
    //
    TDssmBundle::TDssmBundle() {
    }

    inline void AddInputTrait(TDssmBundle::TInputTraits& traits, const TString& fieldName) {
        if (fieldName == "query") {
            traits.UsesQuery = true;
        } else if (fieldName == "query_region") {
            traits.UsesRegion = true;
        } else if (fieldName == "doc_url") {
            // always present
        } else if (fieldName == "doc_path" || fieldName == "doc_host") {
            traits.UsesHostPath = true;
        } else if (fieldName == "doc_title") {
            traits.UsesTitle = true;
        }
    }

    void TDssmBundle::SetBundle(const TDssmModelsConfig& proto, const TFsPath& factorModelsDir) {
        const bool checkExists = proto.GetCheckFilesExistence(); // true means that a missing file is auto-removed from the configuration

        for (const TDssmModelFile& modelFile: proto.get_arr_dssmmodel()) {
            Y_ENSURE(modelFile.HasAlias(), "incorrect TDssmModelsConfig");
            const TString fileName = modelFile.HasFileName() ? modelFile.GetFileName() : modelFile.GetAlias() + ".dssm";
            const TFsPath path = factorModelsDir / fileName;
            if (checkExists && !path.Exists()) {
                DEBUG_LOG << "DSSM model is not available: " << modelFile.GetAlias() << " = " << fileName << Endl;
            } else {
                DEBUG_LOG << "DSSM model is available: " << modelFile.GetAlias() << " = " << path.Basename() << Endl;
                ModelByName.insert({modelFile.GetAlias(), fileName});
            }
        }

        for (const TDssmFactorsBlock& availModel: proto.get_arr_dssmfactors()) {
            Y_ENSURE(!availModel.GetModelAlias().empty(), "incorrect TDssmModelsConfig");
            const TString& modelAlias = availModel.GetModelAlias();
            if (!ModelByName.contains(modelAlias)) {
                DEBUG_LOG << "TDssmFactorsBlock is skipped, because a DSSM model is not there: " << modelAlias << Endl;
                continue;
            }

            TIntrusivePtr<TFactorsBlock> model = MakeIntrusive<TFactorsBlock>();
            model->ModelName = modelAlias;

            TInputTraits& inputs = model->Inputs;
            for (const TString& field: availModel.get_arr_inputfield()) {
                AddInputTrait(inputs, field);
            }

            for (auto&& var: availModel.get_arr_queryembedding()) {
                model->QueryCacheOutputs.emplace_back(std::move(var));
            }

            for (const TDssmFactorsBlock::TFactorDescription& f: availModel.get_arr_factor()) {
                const TString factorName = f.GetFactorName();
                Y_ENSURE(factorName && f.GetOutputVar(), "incorrect TFactorDescription item in TDssmModelsConfig: " << modelAlias << "," << f.ShortDebugString());

                const bool alreadyExists = FactorByName.contains(factorName);
                if (alreadyExists) {
                    DEBUG_LOG << "TDssmFactor " << factorName << " is defined several times, dropped from '" << modelAlias << "'" << Endl;
                    continue;
                }

                model->VarNames.push_back(f.GetOutputVar());

                DEBUG_LOG << "DSSM factor is available: " << factorName << " = " << modelAlias << ":" << f.GetOutputVar() << Endl;

                TFactor fb(model->ModelName, f.GetOutputVar(), model.Get());

                FactorByName.insert(std::make_pair(factorName, fb));
            }

            AvailableModels.insert(std::make_pair(modelAlias, model));
        }
    }

    void TDssmFeature::AddFactor(const TString& factorName, IFeature::TAddFactorFunc adder) {
        adder({factorName, Max<size_t>(), 32, TBitsReader::ftFloat});
        FactorsStaticInfo.insert(factorName);
    }

    void TDssmFeature::InitStaticInfo(TStringBuf modelsPath, IFeature::TAddFactorFunc adder) {
        if (modelsPath.empty()) {
            return;
        }
        for (const auto& factorName : GetFactors(modelsPath)) {
             AddFactor(factorName, adder);
        }
    }

    //
    // TDssmBundle
    //
    TDssmBundle::TFactorBinding TDssmBundle::GetFactor(TStringBuf factorName, bool fastFeaturesOnly) const {
        if (Y_LIKELY(!fastFeaturesOnly)) {
            const auto i = FactorByName.find(factorName);
            const bool found = i != FactorByName.end();
            Y_ENSURE(found, "No binding was found for the SaaS factor named " << factorName);

            return TDssmBundle::TFactorBinding(i->second.ModelName, i->second.VarName);
        } else {
            // Fast Ranking is requested (the first pass in RPF_FASTRANK request, or the single pass in RPF_JUST_FASTRANK request)
            // No precalculated per-document embeddings are implemented right now in SaaS.
            // Hence the DSSM factor is not available on this stage. Will return a canonical value.
            return TDssmBundle::TFactorBinding();
        }
    }

    const TDssmBundle::TModelBinding& TDssmBundle::GetModelMeta(TStringBuf modelName) const {
        const auto i = AvailableModels.find(modelName);
        const bool found = i != AvailableModels.end();
        Y_ENSURE(found && i->second, "No model was found with name=" << modelName);

        return *i->second;
    }

    TString TDssmBundle::GetModelFileName(TStringBuf modelName) const {
        const auto iAlias = ModelByName.find(modelName);
        return iAlias == ModelByName.end() ? (TString)modelName + ".dssm" : iAlias->second;
    }

    TString TDssmBundle::TModelBinding::FormatCacheKey() const {
        TStringStream ss;
        ss << ModelName << "|";
        ss << JoinSeq(",", VarNames) << "|";
        ss << JoinSeq(",", QueryCacheOutputs) << "|";
        return ss.Str();
    }

    //
    // accessors for the unit tests
    //
    void TDssmBundle::DbgAddModel(const TString modelName, TDssmBundle::TInputTraits inputs, TVector<TString>&& queryEmbs) {
        auto modelMeta = MakeIntrusive<TDssmBundle::TFactorsBlock>();
        modelMeta->ModelName = modelName;
        modelMeta->Inputs = inputs;
        modelMeta->QueryCacheOutputs = std::move(queryEmbs);
        const bool inserted = AvailableModels.insert(std::make_pair(modelName, modelMeta)).second;
        Y_ENSURE(inserted);
    }
    void TDssmBundle::DbgAddFactor(const TString& factorName, const TFactorBinding& factor) {
        const auto iter = AvailableModels.find(factor.ModelName);
        Y_ENSURE(iter != AvailableModels.end(), "model " << factor.ModelName << " is not found");
        TDssmBundle::TFactorsBlock* model = iter->second.Get();
        model->VarNames.push_back(factor.VarName);
        const bool inserted = FactorByName.insert(std::make_pair(factorName, TFactor(model->ModelName, factor.VarName, model))).second;
        Y_ENSURE(inserted);
    }

    //
    // TDssmSubmodelsCache
    //
    struct TDssmSubmodelsCache {
        THashMap<TString, NNeuralNetApplier::TModelPtr> Data;

        static TDssmSubmodelsCache* GetInstance(NRTYFactors::TConfig::TFactorModels& cache) {
            const TDssmSubmodelsCache* result = cache.GetOrBuild<TDssmSubmodelsCache>(
                    "TDssmSubmodelsCache",
                    [](TDssmSubmodelsCache&, TStringBuf) {});
            Y_VERIFY(result != nullptr);
            return const_cast<TDssmSubmodelsCache*>(result);
        }

        inline static TString FormatKeyForDocModel(const TString& key) {
            return key + "|doc";
        }

        bool TryGet(const TString& key, NNeuralNetApplier::TModelPtr& querySubmodel, NNeuralNetApplier::TModelPtr& docSubmodel) const {
            Y_VERIFY(key);
            if (!Data.contains(key)) {
                return false;
            }

            const TString keyDoc = FormatKeyForDocModel(key);
            Y_VERIFY(Data.contains(keyDoc));

            querySubmodel = Data.at(key);
            docSubmodel = Data.at(keyDoc);
            Y_ASSERT(querySubmodel && docSubmodel);
            return true;
        }

        void Add(const TString& key, NNeuralNetApplier::TModelPtr& querySubmodel, NNeuralNetApplier::TModelPtr& docSubmodel) {
            Y_VERIFY(key && querySubmodel && docSubmodel);
            Data[key] = querySubmodel;
            Data[FormatKeyForDocModel(key)] = docSubmodel;
        }
    };

    //
    // TDssmFeatureCalcerPlan
    //
    void TDssmFeatureCalcerPlan::DoExtractSubmodels(TDssmFeatureCalcerPlan::TModelDescription& modelDescr) {
        Y_ASSERT(modelDescr.Model);
        Y_ASSERT(!modelDescr.SumBinding.QueryCacheOutputs.empty());

        const auto& b = modelDescr.SumBinding;
        const TSet<TString> embeddingNames(b.QueryCacheOutputs.cbegin(), b.QueryCacheOutputs.cend());
        const TSet<TString> resultNames(b.VarNames.cbegin(), b.VarNames.cend());

        TDssmHelper helper;
        auto fullModel = MakeIntrusive<NNeuralNetApplier::TModel>(*modelDescr.Model);
        helper.SplitFieldExtractor(*fullModel);
        NNeuralNetApplier::TModelPtr querySubmodel = helper.GetQuerySubmodel(*fullModel, modelDescr.ModelAlias, embeddingNames);
        NNeuralNetApplier::TModelPtr docSubmodel = helper.GetDocSubmodel(*fullModel, modelDescr.ModelAlias, resultNames, embeddingNames);

        modelDescr.QuerySubmodel = std::move(querySubmodel);
        modelDescr.DocSubmodel = std::move(docSubmodel);
    }

    bool TDssmFeatureCalcerPlan::FindCachedSubmodels(const NRTYFactors::TConfig& cacheObj, TModelDescription& modelDescr) {
        NRTYFactors::TConfig::TFactorModels& globalCache = cacheObj.GetFactorModels();
        Y_ASSERT(modelDescr.SubmodelCacheKey);
        if (!modelDescr.SubmodelCacheKey)
            return false;

        TDssmSubmodelsCache& cache = *TDssmSubmodelsCache::GetInstance(globalCache);
        return cache.TryGet(modelDescr.SubmodelCacheKey, modelDescr.QuerySubmodel, modelDescr.DocSubmodel);
    }

    void TDssmFeatureCalcerPlan::AddSubmodelsToCache(const NRTYFactors::TConfig& cacheObj, TModelDescription& modelDescr) {
        Y_VERIFY(modelDescr.SubmodelCacheKey && modelDescr.IsSubmodelsCreated());
        NRTYFactors::TConfig::TFactorModels& globalCache = cacheObj.GetFactorModels();
        TDssmSubmodelsCache& cache = *TDssmSubmodelsCache::GetInstance(globalCache);
        cache.Add(modelDescr.SubmodelCacheKey, modelDescr.QuerySubmodel, modelDescr.DocSubmodel);
    }
    //
    // TDssmFeatureCalcer
    //
    namespace {
        class TDssmFeatureCalcer final: public TDssmFeatureCalcerBase {
        private:
            static TString GetPrettyUrl(const IArchiveData* arc, ui32 docId) {
                const TBlob extInfo = arc->GetDocExtInfo(docId)->UncompressBlob();
                TDocDescr descr;
                descr.UseBlob(extInfo.Data(), (ui32)extInfo.Size());
                return descr.get_url();
            }

            static TString GetTitle(const TRTYDynamicFeatureContext& ctx, ui32 docId) {
                const TBlob docText = ctx.AD->GetDocText(docId)->UncompressBlob();
                const TRTYIndexTextData& textOpts = *ctx.TD;
                return NRTYServer::GetArcTitle(docText.AsUnsignedCharPtr(), docId, ctx.Makeup, textOpts.GetTitleZonesMask(), textOpts.GetAutoTitleSentNum());
            }

            static TString GetRegionsChain(const TRequestParams& rp) {
                //TODO(yrum): this is copypasted from reg_chain_dssm_features.cpp
                constexpr TStringBuf EarthReg = "10000";

                TString regChain;
                if (!rp.RearrangeParams.TryGet("all_regions", regChain)) {
                    return TString();
                }

                SubstGlobal(regChain, ',', ' ');

                if (regChain != EarthReg && regChain.EndsWith(EarthReg)) {
                    regChain.resize(regChain.size() - EarthReg.size() - 1);
                }

                return regChain;
            }

            Y_FORCE_INLINE void SaveQueryEmbInputs(const TVector<TString>& names, TVector<TVector<TVector<float>>>& embResults) {
                Y_ASSERT(names.size() == embResults.size());
                Y_ASSERT(std::all_of(embResults.begin(), embResults.end(), [](TVector<TVector<float>>& outputVar) { return outputVar.size() == 1; }));
                for (size_t i = 0; i < names.size(); ++i) {
                    const TString& queryEmbName = names[i];
                    const TVector<float>& queryEmbValue = embResults[i][0];
                    Y_ENSURE(!QueryEmbeddings.contains(queryEmbName), "Query embedding " << queryEmbName << "is provided by more than one DSSM model");
                    Y_ASSERT(queryEmbValue.size() >= 1);
                    //Y_ASSERT(NNeuralNetApplier::IsNormalized(queryEmbValue));

                    QueryEmbeddings[queryEmbName] = new NNeuralNetApplier::TMatrix(1, queryEmbValue.size(), queryEmbValue);
                }
            }

            Y_FORCE_INLINE void FillQueryEmbInputs(NNeuralNetApplier::TEvalContext& eCtx, const TDssmBundle::TModelBinding& sumBinding) const {
                for (size_t i = 0; i < sumBinding.QueryCacheOutputs.size(); ++i) {
                    auto&& queryEmbName = sumBinding.QueryCacheOutputs[i];
                    eCtx[queryEmbName] = QueryEmbeddings.at(queryEmbName);
                }
            }

        public:
            using TDssmFeatureCalcerBase::TDssmFeatureCalcerBase;

        protected:
            void ApplyModelToQuery(const TDssmFeatureCalcerPlan::TModelDescription* m, const TRTYDynamicFeatureContext& ctx) override {
                Y_ENSURE(m && m->Model);
                Y_ENSURE(m->QuerySubmodel);
                Y_ENSURE(!ctx.RP->UserRequest.empty()); // r5857389 made it possible to omit &text= in CGI - but we need that string for DSSM
                const TString& query = ctx.RP->UserRequest;
                const NNeuralNetApplier::TModel* dssmModel = m->QuerySubmodel.Get();
                const TDssmBundle::TModelBinding& sumBinding = m->SumBinding;
                const auto& inputs = sumBinding.Inputs;

                TVector<TString> names{"query"};
                TVector<TString> values{query};
                if (inputs.UsesRegion) {
                    names.push_back("query_region");
                    values.push_back(GetRegionsChain(*ctx.RP));
                }

                TAtomicSharedPtr<NNeuralNetApplier::TSample> mySample = new NNeuralNetApplier::TSample(names, values);
                NNeuralNetApplier::TEvalContext eCtx;
                dssmModel->FillEvalContextFromSample(mySample, eCtx);

                const TVector<TString>& queryEmbNames = m->SumBinding.QueryCacheOutputs;
                TVector<TVector<TVector<float>>> queryEmbResult;
                dssmModel->Apply(eCtx, queryEmbNames, queryEmbResult);
                SaveQueryEmbInputs(queryEmbNames, queryEmbResult);
            }

            void ApplyModel(const TDssmFeatureCalcerPlan::TModelDescription* m, const TRTYDynamicFeatureContext& ctx, ui32 docId) override {
                Y_ASSERT(ctx.AD != nullptr && ctx.TD != nullptr);
                Y_ENSURE(m && m->Model);
                Y_ENSURE(m->SumBinding.QueryCacheOutputs.empty() || !QueryEmbeddings.empty());

                const bool useDocSubmodel = !!m->QuerySubmodel;
                Y_ENSURE(!useDocSubmodel || m->DocSubmodel);

                const NNeuralNetApplier::TModel& dssmModel = !useDocSubmodel ? *m->Model : *m->DocSubmodel;
                const TDssmBundle::TModelBinding& sumBinding = m->SumBinding;
                const auto& inputs = sumBinding.Inputs;

                TVector<TString> names;
                TVector<TString> values;

                if (!useDocSubmodel) {
                    if (inputs.UsesQuery) {
                        names.push_back("query");
                        values.push_back(ctx.RP->UserRequest);
                    }
                    if (inputs.UsesRegion) {
                        names.push_back("query_region");
                        values.push_back(GetRegionsChain(*ctx.RP));
                    }
                }

                const TString url = GetPrettyUrl(ctx.AD, docId);
                names.push_back("doc_url");
                values.push_back(url);
                if (inputs.UsesHostPath) {
                    TString host, path;
                    SplitUrlToHostAndPath(url, host, path);
                    names.push_back("doc_host");
                    names.push_back("doc_path");
                    values.push_back(host);
                    values.push_back(path);
                }

                if (inputs.UsesTitle) {
                    names.push_back("doc_title");
                    values.push_back(GetTitle(ctx, docId));
                }
                //DEBUG_LOG << "ApplyModel\t" << JoinSeq("\t", values) << Endl;

                TAtomicSharedPtr<NNeuralNetApplier::TSample> mySample = new NNeuralNetApplier::TSample(names, values);
                NNeuralNetApplier::TEvalContext eCtx;
                dssmModel.FillEvalContextFromSample(mySample, eCtx);
                FillQueryEmbInputs(eCtx, sumBinding);
                dssmModel.Apply(eCtx, sumBinding.VarNames, Result);
            }

            float GetResult(const ui32 outputLocalId) const override {
                Y_ASSERT(outputLocalId < Result.size());
                return Y_LIKELY(outputLocalId < Result.size()) ? Result[outputLocalId] : 0.5f;
            }

        public:
            TVector<float> Result;
            THashMap<TString, NNeuralNetApplier::TMatrixPtr> QueryEmbeddings;
        };
    }

    //
    // TDssmFeature
    //
    const TDssmBundle& TDssmFeature::GetDssmBundle(const NRTYFactors::TConfig* holder) {
        return *holder->GetFactorModels().GetOrBuild<TDssmBundle>(
            Default<TString>(),
            [](TDssmBundle&, TStringBuf) {
                Y_VERIFY(0, "TDssmBundle is not created (InitModels was not called?)");
            });
    }

    TFsPath TDssmFeature::PathToDssmBundle(TStringBuf modelsPath) {
        TFsPath factorModels = modelsPath;
        Y_ENSURE(factorModels.IsDefined(), "FactorModels config option should be set if DssmFeature is used");
        Y_ENSURE(factorModels.IsDirectory() ||
                (factorModels.IsSymlink() && factorModels.ReadLink().IsDirectory()),
                "FactorModels does not point to a directory"
        );
        factorModels /= "dssm_bundle.pb.txt";
        return factorModels;
    }

    bool TDssmFeature::IsStaticDisable(TStringBuf modelsPath) {
        return modelsPath.empty() || modelsPath == "ZEROFILL";
    }

    TVector<TString> TDssmFeature::GetFactors(TStringBuf modelsPath) {
        if (IsStaticDisable(modelsPath)) {
            return {};
        }

        const TFsPath metaFile = PathToDssmBundle(modelsPath);
        metaFile.CheckExists();

        TDssmModelsConfig protoConfig;
        ParseFromTextFormat(metaFile.GetPath(), protoConfig, EParseFromTextFormatOption::AllowUnknownField);

        TVector<TString> result;
        for (const auto& availModel: protoConfig.get_arr_dssmfactors()) {
            for (const auto& factorBlock: availModel.get_arr_factor()) {
                result.push_back(factorBlock.GetFactorName());
            }
        }
        return result;
    }

    const TDssmBundle& TDssmFeature::GetOrLoadDssmBundle(const NRTYFactors::TConfig* relevCfg, const TString& modelsPath) {
        // GetOrLoadDssmBundle is called only when relev.conf includes some of the rty_dynamic_factors provided by the current plugin.
        // If the FactorsModels variable is set to empty, relev.conf should have no DSSM factors (or we fail in InitModels).
        // If FactorsModels is set to "ZEROFILL", we will fill the provided factors with their canonical values, i.e. zeros (usable for tests, for standalone_indexer etc)

        using namespace NProtoBuf;
        const auto creator = [&](TDssmBundle& b, TStringBuf) {
            const TFsPath factorModelsDir = modelsPath;

            const bool staticDisable = IsStaticDisable(modelsPath);
            TDssmModelsConfig protoConfig;
            if (!staticDisable) {
                const TFsPath metaFile = PathToDssmBundle(modelsPath);
                metaFile.CheckExists();
                ParseFromTextFormat(metaFile.GetPath(), protoConfig, EParseFromTextFormatOption::AllowUnknownField);
            }
            b.SetBundle(protoConfig, factorModelsDir);
            b.SetStaticDisable(staticDisable);
        };
        return *relevCfg->GetFactorModels().GetOrBuild<TDssmBundle>(Default<TString>(), creator);
    }


    const NNeuralNetApplier::TModel* TDssmFeature::GetModel(const NRTYFactors::TConfig* relevCfg, const TString& modelName) {
        NRTYFactors::TConfig::TFactorModels& models = relevCfg->GetFactorModels();
        return models.GetOrBuild<NNeuralNetApplier::TModel>(
            modelName,
            [](NNeuralNetApplier::TModel&, TStringBuf) {
                ythrow yexception() << "The model is expected to be preloaded";
            });
    };

    const NNeuralNetApplier::TModel* TDssmFeature::GetOrLoadModel(const NRTYFactors::TConfig* relevCfg, const TString& modelsPath, const TString& modelName) {
        Y_ENSURE(relevCfg);
        Y_ENSURE(!modelsPath.empty());
        NRTYFactors::TConfig::TFactorModels& models = relevCfg->GetFactorModels();
        const TDssmBundle& bundle = GetDssmBundle(relevCfg);
        Y_ENSURE(!bundle.GetStaticDisable());

        return models.GetOrBuild<NNeuralNetApplier::TModel>(
            modelName,
            [bundle, modelsPath](NNeuralNetApplier::TModel& model, TStringBuf name) {
                const TFsPath dirName = modelsPath ? TFsPath(modelsPath) : TFsPath::Cwd();
                const TFsPath fileName = dirName / bundle.GetModelFileName(name);
                fileName.CheckExists();
                INFO_LOG << "Load DSSM model " << fileName << Endl;
                const TBlob b = TBlob::PrechargedFromFile(fileName);
                model.Load(b);
                model.Init();
            });
    };

    static void ValidateCachedEmbeddings(const TVector<THolder<TDssmFeatureCalcerPlan::TModelDescription>>& bindings) {
        TMap<TString, TString> uniqueEmbNames;
        for (const auto& mdHolder: bindings) {
            for (const TString& embName: mdHolder->SumBinding.QueryCacheOutputs) {
                auto res = uniqueEmbNames.insert(std::make_pair(embName, mdHolder->ModelAlias));
                Y_ENSURE(res.second, "Incorrect SumBinding: embedding '" << res.first->first << "' provided both by '" << mdHolder->ModelAlias << "' and '" << res.first->second << "'");
            }
        }
    }

    THolder<TDssmFeatureCalcerPlan> TDssmFeature::CreatePlan(const TDssmFeature::TFactorsInfo& factorsInfo, const TDssmBundle& modelsInfo, const TDynMapping& factorsList, bool fastFeaturesOnly) {
        THolder<TDssmFeatureCalcerPlan> result = MakeHolder<TDssmFeatureCalcerPlan>();
        const bool staticDisable = modelsInfo.GetStaticDisable();
        if (staticDisable) {
            ERROR_LOG << "TDssmFeatureCalcer is disabled by config, factors will be invalid" << Endl;
        } else {
            for (size_t factorLocalId = 0; factorLocalId < factorsList.size(); factorLocalId++) {
                const auto& f = factorsList[factorLocalId];
                if (f.SourceFactor.empty()) {
                    continue;
                }
                const auto& factorName = f.SourceFactor;
                if (!factorsInfo.contains(factorName)) {
                    continue; // it is not ours
                }

                result->Add(factorLocalId, factorName, modelsInfo, fastFeaturesOnly);
            }
        }
        ValidateCachedEmbeddings(result->GetModels());
        return result;
    };

    void TDssmFeature::PrintDebugInfo(const TDssmFeatureCalcerPlan& plan) {
        for (const auto& item : plan.GetFactors()) {
            const auto& fd = *item;
            Y_VERIFY(fd.ModelDescr != nullptr);
            DEBUG_LOG << "TDssmFeature uses factor " << fd.FactorName << " = " << fd.ModelDescr->ModelAlias << ":" << fd.VarName << Endl;
        }
    }

    bool TDssmFeature::IsFeatureEnabledFor(const TDynMapping& usedFactors) const {
        // if we can handle at least one factor
        return std::any_of(usedFactors.begin(), usedFactors.end(), [this](const auto& checkedFactor)->bool {
            return FactorsStaticInfo.contains(checkedFactor.SourceFactor);
        });
    }

    IFeatureCalcer::TPtr TDssmFeature::CreateCalcer(const TDynMapping& factorsInRequest, const NRTYFactors::TConfig* relevCfg, bool fastFeaturesOnly) {
        if (!IsFeatureEnabledFor(factorsInRequest)) {
            return nullptr;
        }

        THolder<TDssmFeatureCalcerPlan> plan = CreatePlan(FactorsStaticInfo, GetDssmBundle(relevCfg), factorsInRequest, fastFeaturesOnly);
        BindPreloadedModels(*plan, relevCfg);
        return MakeIntrusive<TDssmFeatureCalcer>(std::move(plan), factorsInRequest);
    }

    void TDssmFeatureCalcerBase::Calc(TFactorStorage& storage, const TRTYDynamicFeatureContext& ctx, ui32 docId) {
        const TDssmFeatureCalcerPlan::TModelDescription* m;

        if (IsFirstDoc) {
            for (const auto& mdHolder : Plan->GetModels()) {
                m = mdHolder.Get();
                if (m->SumBinding.QueryCacheOutputs) {
                    ApplyModelToQuery(m, ctx);
                }
            }
            IsFirstDoc = false;
        }

        for (const auto& mdHolder : Plan->GetModels()) {
            m = mdHolder.Get();
            ApplyModel(m, ctx, docId);

            for (const auto& fdHolder : Plan->GetFactors()) {
                const TDssmFeatureCalcerPlan::TFactorDescription* f = fdHolder.Get();
                if (f->ModelDescr != m) {
                    continue;
                }
                const ui32 factorLocalId = f->FactorLocalId;
                Y_ASSERT(factorLocalId < Factors.size());

                const i32 factorIndex = Factors[factorLocalId].BaseIndex;
                Y_ASSERT(factorIndex != -1);
                if (Y_LIKELY(factorIndex != -1)) {
                    const float result = GetResult(f->OutputLocalId);
                    storage[factorIndex] = result;
                }
            }
        }
    }
}
