"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 __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    Object.defineProperty(o, k2, { enumerable: true, get: function() { return m[k]; } });
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DialogueGenerator = exports.BasicSentenceGenerator = void 0;
const ThingTalk = __importStar(require("thingtalk"));
const assert_1 = __importDefault(require("assert"));
const stream_1 = __importDefault(require("stream"));
const I18n = __importStar(require("../i18n"));
const multimap_1 = __importDefault(require("../utils/multimap"));
const random_1 = require("../utils/random");
const Utils = __importStar(require("../utils/entity-utils"));
const ThingTalkUtils = __importStar(require("../utils/thingtalk"));
const generator_1 = __importDefault(require("./generator"));
/**
 * Generate a dataset of single-sentence commands and their associated logical forms.
 */
class BasicSentenceGenerator extends stream_1.default.Readable {
    constructor(options) {
        super({ objectMode: true });
        this._idPrefix = options.idPrefix || '';
        this._locale = options.locale;
        this._langPack = I18n.get(options.locale);
        this._rng = options.rng;
        this._generator = new generator_1.default({
            locale: options.locale,
            timezone: options.timezone,
            templateFiles: options.templateFiles,
            forSide: 'user',
            contextual: false,
            flags: options.flags,
            targetPruningSize: options.targetPruningSize,
            maxDepth: options.maxDepth,
            maxConstants: options.maxConstants || 5,
            debug: options.debug,
            rng: options.rng,
            thingpediaClient: options.thingpediaClient,
            entityAllocator: new ThingTalk.Syntax.SequentialEntityAllocator({}),
            onlyDevices: options.onlyDevices,
            whiteList: options.whiteList
        });
        this._generator.on('progress', (value) => {
            this.emit('progress', value);
        });
        this._initialization = null;
        this._i = 0;
    }
    _read() {
        if (this._initialization === null)
            this._initialization = this._generator.initialize();
        this._initialization.then(() => {
            this._generator.generate([], this._output.bind(this));
            this.emit('progress', this._generator.progress);
            this.push(null);
        }).catch((e) => {
            console.error(e);
            this.emit('error', e);
        });
    }
    _postprocessSentence(derivation, program) {
        let utterance = derivation.toString();
        utterance = utterance.replace(/ +/g, ' ');
        utterance = this._langPack.postprocessSynthetic(utterance, program, this._rng, 'user');
        return utterance;
    }
    _output(depth, derivation) {
        const program = derivation.value.optimize();
        assert_1.default(program !== null); // not-null even after optimize
        let preprocessed = this._postprocessSentence(derivation, program);
        const tokens = preprocessed.split(' ');
        const entities = Utils.makeDummyEntities(preprocessed);
        const tokenized = { tokens, entities };
        const contextEntities = {};
        Utils.renumberEntities(tokenized, contextEntities);
        preprocessed = tokenized.tokens.join(' ');
        let sequence;
        try {
            sequence = ThingTalkUtils.serializePrediction(program, [], tokenized.entities, {
                locale: this._locale
            });
        }
        catch (e) {
            console.error(preprocessed);
            console.error(String(program));
            console.error(sequence);
            console.error(program.prettyprint().trim());
            this.emit('error', e);
            return;
        }
        let id = String(this._i++);
        id = this._idPrefix + depth + '000000000'.substring(0, 9 - id.length) + id;
        const flags = {
            synthetic: true
        };
        this.push({ depth, id, flags, preprocessed, target_code: sequence.join(' ') });
    }
}
exports.BasicSentenceGenerator = BasicSentenceGenerator;
class PartialDialogue {
    constructor(context = null, turns = [], execState = undefined) {
        this.context = context;
        this.turns = turns;
        this.execState = execState;
    }
}
const FACTORS = [50, 75, 75, 100];
/**
 * Generate a minibatch of dialogues.
 *
 * This object is created afresh for every minibatch.
 */
class MinibatchDialogueGenerator {
    constructor(agentGenerator, userGenerator, langPack, simulator, stateValidator, options, minibatchIdx) {
        this._agentGenerator = agentGenerator;
        this._userGenerator = userGenerator;
        this._langPack = langPack;
        this._simulator = simulator;
        this._stateValidator = stateValidator;
        this._minibatchSize = options.minibatchSize;
        this._rng = options.rng;
        this._options = options;
        this._minibatchIdx = minibatchIdx;
        this._turnIdx = 0;
        this._debug = true;
        this._partialDialogues = new random_1.ReservoirSampler(Math.ceil(this._minibatchSize), this._rng);
        this._emptyDialogue = new PartialDialogue();
        this._maybeAddPartialDialog(this._emptyDialogue);
        this._completeDialogues = [];
        for (let turnIdx = 0; turnIdx < options.maxTurns; turnIdx++) {
            const factor = turnIdx < FACTORS.length ? FACTORS[turnIdx] : FACTORS[FACTORS.length - 1];
            this._completeDialogues[turnIdx] = new random_1.ReservoirSampler(Math.ceil(this._minibatchSize * factor), this._rng);
        }
    }
    _maybeAddCompleteDialog(dlg) {
        assert_1.default(dlg.turns.length > 0);
        this._completeDialogues[dlg.turns.length - 1].add(dlg.turns);
    }
    _maybeAddPartialDialog(dlg) {
        const discarded = this._partialDialogues.add(dlg);
        if (discarded !== undefined)
            this._maybeAddCompleteDialog(discarded);
        return discarded !== dlg;
    }
    _postprocessSentence(derivation, program, forTarget) {
        let utterance = derivation.toString();
        utterance = utterance.replace(/ +/g, ' ');
        utterance = this._langPack.postprocessSynthetic(utterance, program, this._rng, forTarget);
        return utterance;
    }
    _generateAgent(partials) {
        const agentTurns = [];
        this._agentGenerator.generate(partials, (depth, derivation) => {
            // derivation.dlg is the PartialDialogue that is being continued
            // derivation.value is the object returned by the root semantic function, with:
            // - state (the thingtalk state)
            // - context (the context info to pass to the semantic function of the user)
            // - tags (context tags to set when generating user sentences)
            // - other properties only relevant to inference time we don't care about
            // set the turn of the agent
            let state = derivation.value.state;
            state = state.optimize();
            assert_1.default(state !== null); // not-null even after optimize
            this._stateValidator.validateAgent(state);
            const utterance = this._postprocessSentence(derivation, state, 'agent');
            const dlg = derivation.context.value;
            assert_1.default(dlg instanceof PartialDialogue);
            const prediction = ThingTalkUtils.computePrediction(dlg.context, state, 'agent');
            const target = prediction.prettyprint();
            agentTurns.push({
                dlg,
                context: derivation.value.context,
                contextPhrases: derivation.value.contextPhrases,
                utterance,
                state,
                target
            });
        });
        return agentTurns;
    }
    _generateUser(continuations, agentTurns) {
        this._userGenerator.generate(agentTurns, (depth, derivation) => {
            // the derivation value for the user is directly the thingtalk user state
            // (unlike the agent)
            let state = derivation.value;
            state = state.optimize();
            assert_1.default(state !== null); // not-null even after optimize
            this._stateValidator.validateUser(state);
            const utterance = this._postprocessSentence(derivation, state, 'user');
            const agentTurn = derivation.context.value;
            const dlg = agentTurn.dlg;
            assert_1.default(dlg instanceof PartialDialogue);
            assert_1.default(dlg === this._emptyDialogue || agentTurn.state);
            const prediction = ThingTalkUtils.computePrediction(agentTurn.state, state, 'user');
            const target = prediction.prettyprint();
            continuations.put(dlg, {
                userState: state,
                newTurn: {
                    context: dlg.context !== null ? dlg.context.prettyprint() : null,
                    agent: agentTurn.utterance,
                    agent_target: agentTurn.target,
                    user: utterance,
                    user_target: target,
                },
            });
        });
    }
    async _continueOneDialogue(dlg, continuations) {
        const ourContinuations = continuations.get(dlg);
        if (ourContinuations.length === 0) {
            // if we have no continuations, mark this dialog as complete
            this._maybeAddCompleteDialog(dlg);
        }
        else {
            for (const { userState, newTurn } of ourContinuations) {
                const newDialogue = new PartialDialogue(dlg.context, dlg.turns.concat([newTurn]), dlg.execState);
                if (this._maybeAddPartialDialog(newDialogue)) {
                    const { newDialogueState, newExecutorState } = await this._simulator.execute(userState, dlg.execState);
                    newDialogue.context = newDialogueState;
                    newDialogue.execState = newExecutorState;
                }
            }
        }
    }
    async nextTurn() {
        this._turnIdx++;
        const start = Date.now();
        if (this._debug)
            console.log(`Minibatch ${this._minibatchIdx}, turn ${this._turnIdx}`);
        const partials = this._partialDialogues.sampled;
        this._partialDialogues.reset();
        const continuations = new multimap_1.default();
        let agentTurns = [];
        if (this._turnIdx > 1) { // turnIdx is 1-based because we incremented it already
            agentTurns = this._generateAgent(partials);
        }
        else {
            agentTurns = [{
                    dlg: this._emptyDialogue,
                    context: null,
                    contextPhrases: [],
                    utterance: '',
                    state: null,
                    target: null,
                }];
        }
        this._generateUser(continuations, agentTurns);
        for (const dlg of partials)
            await this._continueOneDialogue(dlg, continuations);
        const end = Date.now();
        if (this._debug) {
            console.log(`Produced ${this._partialDialogues.counter} partial dialogs this turn`);
            console.log(`Turn took ${Math.round((end - start) / 1000)} seconds`);
        }
    }
    *complete() {
        for (const dlg of this._partialDialogues)
            this._maybeAddCompleteDialog(dlg);
        //assert(this._completeDialogs.length > 0);
        for (let turnIdx = 0; turnIdx < this._options.maxTurns; turnIdx++) {
            for (const dialogue of this._completeDialogues[turnIdx])
                yield dialogue;
        }
    }
}
/**
 * Generate a dataset of multi-turn dialogues.
 */
class DialogueGenerator extends stream_1.default.Readable {
    constructor(options) {
        super({ objectMode: true });
        this._i = 0;
        this._numMinibatches = options.numMinibatches;
        this._options = options;
        this._idPrefix = options.idPrefix || '';
        this._debug = options.debug;
        this._langPack = I18n.get(options.locale);
        const agentOptions = {
            locale: options.locale,
            timezone: options.timezone,
            templateFiles: options.templateFiles,
            rootSymbol: '$agent',
            forSide: 'agent',
            contextual: true,
            flags: options.flags,
            targetPruningSize: options.targetPruningSize,
            maxDepth: options.maxDepth,
            maxConstants: options.maxConstants || 5,
            debug: options.debug,
            rng: options.rng,
            thingpediaClient: options.thingpediaClient,
            entityAllocator: new ThingTalk.Syntax.SequentialEntityAllocator({}),
            onlyDevices: options.onlyDevices,
            whiteList: options.whiteList,
            contextInitializer: (partialDialogue, functionTable, contextTable) => {
                return functionTable.context(partialDialogue.context, contextTable);
            }
        };
        this._agentGenerator = new generator_1.default(agentOptions);
        const userOptions = {
            locale: options.locale,
            timezone: options.timezone,
            templateFiles: options.templateFiles,
            rootSymbol: '$user',
            forSide: 'user',
            contextual: true,
            flags: options.flags,
            targetPruningSize: options.targetPruningSize,
            maxDepth: options.maxDepth,
            maxConstants: options.maxConstants || 5,
            debug: options.debug,
            rng: options.rng,
            thingpediaClient: options.thingpediaClient,
            entityAllocator: new ThingTalk.Syntax.SequentialEntityAllocator({}),
            onlyDevices: options.onlyDevices,
            whiteList: options.whiteList,
            contextInitializer: (agentTurn, functionTable, contextTable) => {
                if (agentTurn.context === null) {
                    // first turn
                    return functionTable.context(null, contextTable);
                }
                else {
                    return agentTurn.contextPhrases;
                }
            }
        };
        this._userGenerator = new generator_1.default(userOptions);
        this._stateValidator = ThingTalkUtils.createStateValidator(options.policyFile);
        this._initialized = false;
        this._simulator = ThingTalkUtils.createSimulator({
            locale: options.locale,
            timezone: options.timezone,
            thingpediaClient: options.thingpediaClient,
            database: options.database,
            rng: options.rng,
            interactive: false
        });
        this._minibatchIdx = 0;
    }
    async _generateMinibatch() {
        if (!this._initialized) {
            await this._userGenerator.initialize();
            await this._agentGenerator.initialize();
            await this._stateValidator.load();
            this._initialized = true;
        }
        const start = Date.now();
        let counter = 0;
        try {
            const generator = new MinibatchDialogueGenerator(this._agentGenerator, this._userGenerator, this._langPack, this._simulator, this._stateValidator, this._options, this._minibatchIdx++);
            for (let turn = 0; turn < this._options.maxTurns; turn++)
                await generator.nextTurn();
            for (const turns of generator.complete()) {
                const dlg = {
                    id: this._idPrefix + '' + this._i++,
                    turns
                };
                this.push(dlg);
                counter++;
            }
        }
        finally {
            const end = Date.now();
            if (this._debug)
                console.log(`Minibatch took ${Math.round((end - start) / 1000)} seconds and produced ${counter} dialogues`);
        }
    }
    _read() {
        if (this._minibatchIdx >= this._numMinibatches) {
            this.push(null);
            return;
        }
        this._generateMinibatch().catch((e) => {
            this.emit('error', e);
        });
    }
}
exports.DialogueGenerator = DialogueGenerator;
//# sourceMappingURL=batch.js.map