#include <dict/word2vec/model/model.h>
#include <robot/library/yt/static/command.h>
#include <wmconsole/version3/searchqueries-mr/tools/host2vec-model/protos/embedding.pb.h>

#include "utils.h"

namespace NWebmaster {
namespace NHost2Vec {

void UpdateFilteredModel(NYT::IClientBasePtr client, const TString &fltTable, const TString &fltFileWords, const TString &fltFileVectors) {
    NYT::ITransactionPtr tx = client->StartTransaction();
    auto fltTableObj = NJupiter::TTable<NHost2Vec::NProto::TEmbedding>(tx, fltTable).SelectFields({"host", "vector"});
    auto reader = fltTableObj.GetReader();
    const size_t wordsCount = fltTableObj.GetRecordsCount();
    const size_t vectorSize = reader->GetRow().Getvector().size() / sizeof(NWord2Vec::TCoordinate);
    NYT::IFileWriterPtr outWords = tx->CreateFileWriter(fltFileWords);
    NYT::IFileWriterPtr outVectors = tx->CreateFileWriter(fltFileVectors);
    *outWords << wordsCount << ' ' << vectorSize << Endl;
    for (; reader->IsValid(); reader->Next()) {
        const auto &row = reader->GetRow();
        *outWords << row.Gethost() << Endl;
        outVectors->Write(row.Getvector().data(), row.Getvector().size());
    }
    outWords->Finish();
    outVectors->Finish();
    tx->Commit();
}

} //namespace NHost2Vec
} //namespace NWebmaster
