"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 });
const stream_1 = __importDefault(require("stream"));
const random_1 = require("../../utils/random");
const sample_utils_1 = require("./sample-utils");
const I18n = __importStar(require("../../i18n"));
const ThingTalkUtils = __importStar(require("../../utils/thingtalk"));
class UnassignableEntity extends Error {
}
class SkippedEntity extends Error {
}
class SentenceProcessor {
    constructor(tpClient, schemaRetriever, constants, detokenizer, options, input) {
        this._tpClient = tpClient;
        this._schemaRetriever = schemaRetriever;
        this._constants = constants;
        this._allContexts = options.contexts || new Map;
        this._detokenizer = detokenizer;
        this._functionBlackList = options.functionBlackList;
        this._deviceBlackList = options.deviceBlackList;
        this._functionWhiteList = options.functionWhiteList;
        this._deviceWhiteList = options.deviceWhiteList;
        this._rng = options.rng;
        this._compoundOnly = options.compoundOnly;
        this._debug = options.debug;
        this._locale = options.locale;
        this._id = input.id;
        this._context = input.context || null;
        this._sentence = input.preprocessed;
        this._targetCode = String(input.target_code);
        this._assignedEntities = {};
        this._usedValues = new Set;
    }
    _entityRetriever(entity, param, functionname, unit, options) {
        if (this._assignedEntities[entity])
            return this._assignedEntities[entity].value;
        const underscoreindex = entity.lastIndexOf('_');
        const entitytype = entity.substring(0, underscoreindex);
        // special handling for PATH_NAME (HACK)
        if (!options.forContext && entitytype === 'PATH_NAME' && (param === 'repo_name' || param === 'folder_name'))
            throw new SkippedEntity;
        // FIXME: makeLookupKeys always return non-Array types, there might be cases we need array types
        const ttType = sample_utils_1.entityTypeToTTType(entitytype, unit);
        const keys = sample_utils_1.makeLookupKeys(functionname, param, ttType);
        let choices;
        for (const key of keys) {
            if (this._constants[key]) {
                choices = this._constants[key];
                break;
            }
        }
        if (!choices)
            throw new Error('unrecognized entity type ' + entitytype);
        // special handling for NUMBER followed by a unit (measure) because the unit is chosen
        // by the construct template
        // (this is not needed for DURATION/MEASURE_* entities, which carry their own unit instead)
        if (ttType.isMeasure && entitytype === 'NUMBER')
            choices = choices.filter((c) => c.unit === unit);
        //let index = parseInt(entity.substring(underscoreindex+1));
        if (choices.length > 0) {
            for (let i = 0; i < 10; i++) {
                const choice = random_1.uniform(choices, this._rng);
                let value = choice.value;
                // checks if number are in order, disable for schema.org
                /*if (entitytype === 'NUMBER' && this._assignedEntities['NUMBER_' + (index-1)] && this._assignedEntities['NUMBER_' + (index-1)].value >= value)
                    continue;
                if (entitytype === 'NUMBER' && this._assignedEntities['NUMBER_' + (index+1)] && this._assignedEntities['NUMBER_' + (index+1)].value <= value)
                    continue;*/
                if (!this._usedValues.has(choice.key === null && choice.display ? choice.display : choice.key)) {
                    let display;
                    // ignore display for measure/NUMBER, and just the value for measure/DURATION
                    if (ttType.isMeasure) {
                        if (entitytype === 'NUMBER') {
                            display = value.toLocaleString(this._locale);
                        }
                        else if (entitytype === 'DURATION' || entitytype.startsWith('MEASURE_')) {
                            value = { value: value, unit: choice.unit };
                            if (choice.display)
                                display = choice.display;
                            else
                                display = value.value.toLocaleString(this._locale) + ' ' + choice.unit;
                        }
                        else {
                            throw new TypeError('???');
                        }
                    }
                    else {
                        display = choice.display;
                    }
                    this._assignedEntities[entity] = { value, display };
                    this._usedValues.add(choice.key === null && choice.display ? choice.display : choice.key);
                    return value;
                }
            }
        }
        throw new UnassignableEntity(`Run out of values for ${entity} (unit ${unit}, param name ${param})`);
    }
    _filteredByHeuristics(code) {
        // catch gmail reply and twitter retweet
        let hasGmailInbox = false;
        for (let i = 0; i < code.length; i++) {
            const token = code[i];
            if (token === '@com.gmail.inbox') {
                hasGmailInbox = true;
                continue;
            }
            if (hasGmailInbox &&
                (token === '@com.gmail.send_email' || token === '@com.gmail.send_picture'))
                return true;
        }
        let hasTweetInbox = false;
        for (let i = 0; i < code.length; i++) {
            const token = code[i];
            if (token === '@com.twitter.search' || token === '@com.twitter.home_timeline' || token === '@com.twitter.my_tweets') {
                hasTweetInbox = true;
                continue;
            }
            if (hasTweetInbox &&
                (token === '@com.twitter.post' || token === '@com.twitter.post_picture'))
                return true;
        }
        return false;
    }
    async process() {
        const code = this._targetCode;
        const tokens = code.split(' ');
        if (this._filteredByHeuristics(tokens))
            return null;
        let num_pp = 0;
        let num_filters = 0;
        const functions = [];
        for (let i = 0; i < tokens.length; i++) {
            const token = tokens[i];
            if (token.startsWith('@') && token !== '@org.thingpedia.builtin.thingengine.builtin.say') {
                if (this._functionBlackList && this._functionBlackList.has(token))
                    return null;
                if (this._functionWhiteList && !this._functionWhiteList.has(token))
                    return null;
                const dot = token.lastIndexOf('.');
                const deviceName = token.substring(0, dot) + '.*';
                if (this._deviceBlackList && this._deviceBlackList.has(deviceName))
                    return null;
                if (this._deviceWhiteList && !this._deviceWhiteList.has(deviceName))
                    return null;
                functions.push(token);
            }
            if (token.startsWith('param:')
                && i < tokens.length - 2 &&
                tokens[i + 1] === '=' &&
                tokens[i + 2].startsWith('param:'))
                num_pp++;
            if (token.startsWith('param:')
                && i < tokens.length - 1
                && ['==', '>=', '<=', '=~', '!=', 'contains', 'in_array', 'starts_with', 'ends_with'].indexOf(tokens[i + 1]) >= 0
                && i >= 2
                && !tokens[i - 2].startsWith('param:'))
                num_filters++;
        }
        const function_signature = functions.join('+');
        const num_functions = functions.length;
        if (this._compoundOnly && num_functions < 2)
            return null;
        let prim_type = 'compound';
        if (num_functions < 2) {
            if (tokens[0] === 'now' && tokens[tokens.length - 1] !== 'notify')
                prim_type = 'action';
            else
                prim_type = 'query';
        }
        let context = null;
        if (this._context !== null) {
            const contexts = this._allContexts.get(this._context);
            if (!contexts || contexts.length === 0) {
                if (this._debug)
                    console.log(`Skipped ${this._id} because the context is not known`);
                return null;
            }
            try {
                const entityResolver = ((entity, param, functionname, unit) => this._entityRetriever(entity, param, functionname, unit, { forContext: true }));
                context = await ThingTalkUtils.parsePrediction(this._context.split(' '), entityResolver, {
                    thingpediaClient: this._tpClient,
                    schemaRetriever: this._schemaRetriever
                }, true);
            }
            catch (e) {
                if (e instanceof SkippedEntity)
                    return null;
                if (!(e instanceof UnassignableEntity))
                    throw e;
                if (this._debug)
                    console.log(`Skipped ${this._id} due to context: ${e.message}`);
                return null;
            }
        }
        let program;
        try {
            const entityResolver = ((entity, param, functionname, unit) => this._entityRetriever(entity, param, functionname, unit, { forContext: false }));
            program = await ThingTalkUtils.parsePrediction(tokens, entityResolver, {
                thingpediaClient: this._tpClient,
                schemaRetriever: this._schemaRetriever
            }, true);
        }
        catch (e) {
            if (e instanceof SkippedEntity)
                return null;
            if (!(e instanceof UnassignableEntity))
                throw e;
            if (this._debug)
                console.log(`Skipped ${this._id}: ${e.message}`);
            return null;
        }
        let sentence = '';
        let prevtoken = null;
        let num_entities = 0;
        for (let token of this._sentence.split(' ')) {
            // replace entities and undo penn tree bank tokenization
            if (/^[A-Z]/.test(token)) { // entity
                num_entities++;
                if (!this._assignedEntities[token]) {
                    console.error(this._sentence);
                    throw new Error(`Missing entity ${token} (present in the sentence, not in the code)`);
                }
                token = this._assignedEntities[token].display;
            }
            sentence = this._detokenizer(sentence, prevtoken, token);
            prevtoken = token;
        }
        let context_utterance = '';
        const assistant_action = ''; // FIXME
        if (context !== null) {
            let prevtoken = null;
            const choices = this._allContexts.get(this._context);
            for (let token of random_1.uniform(choices, this._rng).split(' ')) {
                // replace entities and undo penn tree bank tokenization
                if (/^[A-Z]/.test(token)) { // entity
                    if (!this._assignedEntities[token]) {
                        console.log(this._context, this._targetCode, this._assignedEntities);
                        throw new Error(`Missing entity ${token} (present in the context sentence, not in the context code)`);
                    }
                    token = this._assignedEntities[token].display;
                }
                context_utterance = this._detokenizer(context_utterance, prevtoken, token);
                prevtoken = token;
            }
        }
        // remove flags
        const id = this._id.replace(/^R?P?S/, '');
        const depth = parseInt(id[0]);
        const sentence_length = this._sentence.length;
        const obj = {
            id: this._id,
            utterance: sentence,
            target_code: program.prettyprint(),
            depth: depth,
            sentence_length,
            num_functions,
            num_entities,
            num_pp,
            num_filters,
            prim_type,
            function_signature,
        };
        if (context !== null) {
            obj.context = context.prettyprint();
            obj.context_utterance = context_utterance;
            obj.assistant_action = assistant_action;
        }
        return obj;
    }
}
function remove_units(code) {
    return code.replace(/unit:\S+/g, '');
}
const SAMPLING_STRATEGIES = {
    bySignature: {
        init() {
            return new Map;
        },
        iterate(bags) {
            return bags;
        },
        collect(bags, input, result) {
            const functionsig = result.function_signature;
            if (!bags.has(functionsig))
                bags.set(functionsig, []);
            bags.get(functionsig).push(result);
        },
        sample(sig, choices, options) {
            const signature = sig.split('+');
            let chosen = [];
            if (signature.length === 1) {
                if (choices[0].prim_type === 'query') {
                    if (options.debug)
                        console.log('primitive query: ' + signature[0] + ' ' + choices.length);
                    chosen = random_1.choose(choices, 50, options.rng);
                }
                else {
                    if (options.debug)
                        console.log('primitive action: ' + signature[0] + ' ' + choices.length);
                    chosen = random_1.choose(choices, 20, options.rng);
                }
            }
            else if (signature.length === 2) {
                if (signature[0] === signature[1])
                    return [];
                if (signature.every((sig) => options.functionHighValueList && options.functionHighValueList.has(sig))) {
                    if (options.debug)
                        console.log('high value compound: ' + signature.join('+') + ' ' + choices.length);
                    chosen = random_1.choose(choices, 10, options.rng);
                }
                else if (signature.some((sig) => options.functionHighValueList && options.functionHighValueList.has(sig))) {
                    if (options.debug)
                        console.log('mid value compound: ' + signature.join('+') + ' ' + choices.length);
                    chosen = random_1.choose(choices, 1, options.rng);
                }
                else {
                    if (options.debug)
                        console.log('low value compound: ' + signature.join('+') + ' ' + choices.length);
                    if (random_1.coin(0.05, options.rng))
                        chosen = [random_1.uniform(choices, options.rng)];
                }
            }
            else if (signature.length === 3) {
                if (random_1.coin(0.5, options.rng))
                    chosen = [random_1.uniform(choices, options.rng)];
            }
            if (options.debug)
                console.log('produced for ' + signature.join('+') + ' : ' + chosen.length);
            return chosen;
        }
    },
    byCode: {
        init() {
            return new Map;
        },
        iterate(bags) {
            return bags;
        },
        collect(bags, input, result) {
            const unified = remove_units(String(input.target_code));
            if (!bags.has(unified))
                bags.set(unified, []);
            bags.get(unified).push(result);
        },
        sample(key, choices, options) {
            return random_1.choose(choices, 3, options.rng);
        }
    },
    bySentence: {
        init() {
            return {
                s: 0,
                sentences: [],
                bags: new Map
            };
        },
        iterate(state) {
            return state.bags;
        },
        collect(state, input, result, options) {
            const sentences = state.sentences;
            const bags = state.bags;
            const newSentence = input.preprocessed;
            let bag = bags.get(newSentence);
            if (!bag) {
                const newBag = [];
                newBag.n = 0;
                bag = newBag;
                if (sentences.length < 20000) {
                    sentences.push(newSentence);
                    bags.set(newSentence, bag);
                }
                else if (random_1.coin(20000 / state.s, options.rng)) {
                    const i = Math.floor(options.rng() * 20000);
                    const toDelete = sentences[i];
                    bags.delete(toDelete);
                    sentences[i] = newSentence;
                }
                else {
                    // drop this sentence on the floor
                    state.s++;
                    return;
                }
                bags.set(newSentence, newBag);
                state.s++;
            }
            const bag2 = bag;
            if (bag2.length < 3) {
                bag2.push(result);
            }
            else if (random_1.coin(3 / bag2.n, options.rng)) {
                const i = Math.floor(options.rng() * 3);
                bag2[i] = result;
            }
            bag2.n++;
        },
        sample(key, choices, options) {
            return choices; //choose(choices, 3, options.rng);
        }
    }
};
class SentenceSampler extends stream_1.default.Transform {
    constructor(tpClient, schemaRetriever, constants, options) {
        super({
            readableObjectMode: true,
            writableObjectMode: true,
        });
        this._constants = constants;
        this._tpClient = tpClient;
        this._schemaRetriever = schemaRetriever;
        this._options = options;
        if (!this._options.functionHighValueList)
            this._options.functionHighValueList = new Set;
        this._detokenizer = I18n.get(options.locale).detokenize;
        this._samplingStrategy = SAMPLING_STRATEGIES[options.samplingStrategy || 'bySignature'];
        if (!this._samplingStrategy)
            throw new TypeError(`Invalid sampling strategy ${options.samplingStrategy}`);
        this._samplingState = this._samplingStrategy.init();
    }
    async _run(input) {
        const processor = new SentenceProcessor(this._tpClient, this._schemaRetriever, this._constants, this._detokenizer, this._options, input);
        const result = await processor.process();
        if (result === null)
            return;
        this._samplingStrategy.collect(this._samplingState, input, result, this._options);
    }
    _transform(input, encoding, callback) {
        this._run(input).then(() => callback(), callback);
    }
    _flush(callback) {
        for (const [sig, choices] of this._samplingStrategy.iterate(this._samplingState)) {
            if (choices.length === 0)
                continue;
            const chosen = this._samplingStrategy.sample(sig, choices, this._options);
            for (const c of chosen) {
                const ret = {
                    id: c.id
                };
                if (c.context) {
                    ret.context = c.context;
                    ret.context_utterance = c.context_utterance;
                    ret.assistant_action = c.assistant_action;
                }
                ret.utterance = c.utterance;
                ret.target_code = c.target_code;
                this.push(ret);
            }
        }
        callback();
    }
}
exports.default = SentenceSampler;
//# sourceMappingURL=sampler.js.map