"use strict";
// -*- mode: typescript; indent-tabs-mode: nil; js-basic-offset: 4 -*-
//
// This file is part of Genie
//
// Copyright 2020 The Board of Trustees of the Leland Stanford Junior University
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Author: Giovanni Campagna <gcampagn@cs.stanford.edu>
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.SPLIT_STRATEGIES = void 0;
const stream_1 = __importDefault(require("stream"));
const random_1 = require("../utils/random");
const requoting_1 = require("./requoting");
function filterForDevices(set, code) {
    for (const token of code.split(' ')) {
        if (token.startsWith('@')) {
            const split = token.substring(1).split('.');
            const kind = split.slice(0, split.length - 1).join('.');
            if (!set.has(kind))
                return false;
        }
    }
    return true;
}
function join(iterable) {
    return Array.from(iterable).join(' ');
}
exports.SPLIT_STRATEGIES = {
    id: (id) => id,
    'raw-sentence': (id, sentence) => sentence,
    sentence: (id, sentence, program) => join(requoting_1.requoteSentence(id, sentence, program)),
    spotify: (id, sentence, program) => join(requoting_1.requoteSentence(id, sentence, program)),
    aggregate: (id, sentence, program) => {
        if (program.split(' ').indexOf('aggregate') < 0)
            return ':train';
        return join(requoting_1.requoteSentence(id, sentence, program));
    },
    program: (id, sentence, program) => join(requoting_1.requoteProgram(program)),
    combination: (id, sentence, program) => {
        const functions = Array.from(requoting_1.getFunctions(program));
        if (functions.length <= 1)
            return ':train';
        else
            return functions.join(' ');
    },
    'context-and-program': (id, sentence, program, context) => context + ' ' + program
};
class DatasetSplitter extends stream_1.default.Writable {
    constructor(options) {
        super({ objectMode: true });
        this._rng = options.rng;
        this._evalOnSynthetic = options.evalOnSynthetic;
        this._useEvalFlag = options.useEvalFlag;
        this._train = options.train;
        this._eval = options.eval;
        this._test = options.test;
        this._evalProbability = options.evalProbability;
        this._forDevices = new Set(options.forDevices || []);
        this._splitStrategy = exports.SPLIT_STRATEGIES[options.splitStrategy || 'sentence'];
        this._dedupedevtestMakeKey = undefined;
        if (options.splitStrategy === 'sentence' || options.splitStrategy === 'raw-sentence')
            this._dedupedevtestMakeKey = (splitKey) => splitKey;
        else if (options.splitStrategy === 'program' || options.splitStrategy === 'combination')
            this._dedupedevtestMakeKey = (splitKey, id, sentence, program) => exports.SPLIT_STRATEGIES.sentence(id, sentence, program);
        this._devtestset = new Set;
        this._dedupeddevtestset = new Set;
        this._trainset = new Set;
        this._trainset.add(':train');
    }
    _final(callback) {
        this._train.end();
        this._eval.end();
        if (this._test)
            this._test.end();
        callback();
    }
    /**
      Check if this example can potentially be used for evaluation.

      If this method returns true, the example is a candidate for sampling
      in the evaluation sets (validation/test).

      Synthetic (S) and augmented (P) sentences are excluded if `evalOnSynthetic`
      is false (default),

      If `useEvalFlag` was passed as option, this method considers the `eval` flag
      (E flag in TSV format), otherwise all other sentences are potentially included
      in the evaluation sets.
     */
    _isFlaggedForEval(flags) {
        if (!this._evalOnSynthetic && (flags.synthetic || flags.augmented))
            return false;
        if (this._useEvalFlag)
            return !!flags.eval;
        else
            return true;
    }
    async _doWriteLower(stream, row) {
        await new Promise((resolve, reject) => {
            stream.write(row, (err) => {
                if (err)
                    reject(err);
                else
                    resolve();
            });
        });
    }
    async _handleOne(row) {
        if (this._forDevices.size > 0 && !filterForDevices(this._forDevices, row.target_code))
            return;
        const splitKey = this._splitStrategy(row.id, row.preprocessed, row.target_code, row.context);
        const flags = row.flags || {};
        //console.log(flags, splitKey, this._devtestset);
        if (!this._isFlaggedForEval(flags)) {
            if (splitKey !== undefined && this._devtestset.has(splitKey))
                return;
            await this._doWriteLower(this._train, row);
        }
        else {
            if (splitKey !== undefined && this._devtestset.has(splitKey)) {
                if (this._dedupedevtestMakeKey) {
                    const dedupeKey = this._dedupedevtestMakeKey(splitKey, row.id, row.preprocessed, row.target_code);
                    if (this._dedupeddevtestset.has(dedupeKey))
                        return;
                    this._dedupeddevtestset.add(dedupeKey);
                }
                if (this._test && random_1.coin(0.5, this._rng))
                    await this._doWriteLower(this._test, row);
                else
                    await this._doWriteLower(this._eval, row);
            }
            else if (splitKey !== undefined && this._trainset.has(splitKey)) {
                await this._doWriteLower(this._train, row);
            }
            else if (random_1.coin(this._evalProbability, this._rng)) {
                if (this._dedupedevtestMakeKey) {
                    const dedupeKey = this._dedupedevtestMakeKey(splitKey, row.id, row.preprocessed, row.target_code);
                    if (this._dedupeddevtestset.has(dedupeKey))
                        return;
                    this._dedupeddevtestset.add(dedupeKey);
                }
                if (this._test && random_1.coin(0.5, this._rng))
                    await this._doWriteLower(this._test, row);
                else
                    await this._doWriteLower(this._eval, row);
                if (splitKey !== undefined)
                    this._devtestset.add(splitKey);
            }
            else {
                await this._doWriteLower(this._train, row);
                if (splitKey !== undefined)
                    this._trainset.add(splitKey);
            }
        }
    }
    async _handleMany(row) {
        if (typeof row.target_code === 'string') {
            await this._handleOne({
                id: row.id,
                flags: row.flags,
                preprocessed: row.preprocessed,
                target_code: row.target_code,
                context: row.context
            });
            return;
        }
        for (const code of row.target_code) {
            await this._handleOne({
                id: row.id,
                flags: row.flags,
                preprocessed: row.preprocessed,
                target_code: code,
                context: row.context
            });
        }
    }
    _write(row, encoding, callback) {
        this._handleMany(row).then(() => callback(), (err) => callback(err));
    }
}
exports.default = DatasetSplitter;
//# sourceMappingURL=splitter.js.map