"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.CollectDialogueStatistics = exports.DialogueEvaluatorStream = exports.KEYS = void 0;
const stream_1 = __importDefault(require("stream"));
const deep_equal_1 = __importDefault(require("deep-equal"));
const ThingTalk = __importStar(require("thingtalk"));
const thingtalk_1 = require("thingtalk");
const assert_1 = __importDefault(require("assert"));
const Utils = __importStar(require("../../utils/misc-utils"));
const I18n = __importStar(require("../../i18n"));
const entity_finder_1 = require("../../dialogue-agent/entity-linking/entity-finder");
const eval_utils_1 = require("./eval_utils");
const ThingTalkUtils = __importStar(require("../../utils/thingtalk"));
const MINIBATCH_SIZE = 100;
class DialogueEvaluatorStream extends stream_1.default.Transform {
    constructor(parser, options) {
        super({ objectMode: true });
        this._parser = parser;
        this._tpClient = options.thingpediaClient;
        this._tokenizer = I18n.get(options.locale).getTokenizer();
        this._options = options;
        this._locale = options.locale;
        this._debug = !!options.debug;
        this._tokenized = options.tokenized;
        this._database = options.database;
        this._cachedEntityMatches = new Map;
        this._minibatch = [];
    }
    async _preprocess(sentence, contextEntities) {
        let tokenized;
        if (this._tokenized) {
            const tokens = sentence.split(' ');
            const entities = Utils.makeDummyEntities(sentence);
            tokenized = { tokens, entities };
        }
        else {
            tokenized = this._tokenizer.tokenize(sentence);
        }
        Utils.renumberEntities(tokenized, contextEntities);
        return tokenized;
    }
    _isWellKnownEntity(entityType) {
        switch (entityType) {
            case 'tt:username':
            case 'tt:hashtag':
            case 'tt:picture':
            case 'tt:url':
            case 'tt:email_address':
            case 'tt:phone_number':
            case 'tt:path_name':
            case 'tt:device':
            case 'tt:function':
                return true;
            default:
                return false;
        }
    }
    _tokenizeSlot(value) {
        return this._tokenizer.tokenize(value).rawTokens.join(' ');
    }
    async _resolveEntity(value) {
        if (this._isWellKnownEntity(value.type)) {
            assert_1.default(value.value);
            return { value: value.value, name: value.display || '', canonical: value.value };
        }
        const searchKey = value.display || value.value;
        assert_1.default(searchKey);
        const cacheKey = value.type + '/' + value.value + '/' + searchKey;
        let resolved = this._cachedEntityMatches.get(cacheKey);
        if (resolved)
            return resolved;
        if (this._database && this._database.has(value.type)) {
            // resolve as ID entity from the database (simulate issuing a query for it)
            const ids = this._database.get(value.type).map((entry) => {
                const id = entry.id;
                return {
                    value: id.value,
                    name: id.display,
                    canonical: id.display.toLowerCase()
                };
            });
            if (value.value) {
                for (const id of ids) {
                    if (id.value === value.value) {
                        resolved = id;
                        break;
                    }
                }
            }
            if (!resolved)
                resolved = entity_finder_1.getBestEntityMatch(searchKey, value.type, ids);
            this._cachedEntityMatches.set(cacheKey, resolved);
            return resolved;
        }
        // resolve as regular Thingpedia entity
        const candidates = await this._tpClient.lookupEntity(value.type, searchKey);
        resolved = entity_finder_1.getBestEntityMatch(searchKey, value.type, candidates.data);
        this._cachedEntityMatches.set(cacheKey, resolved);
        return resolved;
    }
    async _valueToSlot(value) {
        // HACK
        if (value instanceof thingtalk_1.Ast.ComputationValue)
            return this._valueToSlot(value.operands[0]);
        if (value instanceof thingtalk_1.Ast.EntityValue) {
            const resolved = await this._resolveEntity(value);
            if (resolved)
                return resolved.canonical;
            return this._tokenizeSlot(value.display || '');
        }
        if (value instanceof thingtalk_1.Ast.BooleanValue)
            return value.value ? 'yes' : 'no';
        if (value instanceof thingtalk_1.Ast.LocationValue) {
            if (value.value instanceof thingtalk_1.Ast.RelativeLocation)
                return value.value.relativeTag;
            if (value.value instanceof thingtalk_1.Ast.AbsoluteLocation)
                return this._tokenizeSlot(value.value.display || '');
            // unresolved
            assert_1.default(value.value instanceof thingtalk_1.Ast.UnresolvedLocation);
            return this._tokenizeSlot(value.value.name);
        }
        if (value instanceof thingtalk_1.Ast.ContextRefValue)
            return 'context-' + value.name;
        if (value instanceof thingtalk_1.Ast.StringValue) {
            // "tokenize" the value, because the prediction will also be tokenized
            return this._tokenizeSlot(value.toJS());
        }
        // everything else (time, currency, number, enum), use JS value
        return String(value.toJS()).toLowerCase();
    }
    async _extractSlots(state) {
        const slotValues = {};
        let currentDomain;
        function nameToSlot(domain, name) {
            if (name === 'id' || name === domain)
                return domain + '-name';
            const slotKey = domain + '-' + name.replace(/_/g, '-');
            return slotKey;
        }
        // note: this function relies on the precise visit order, in which an invocation
        // is visited before the boolean expressions that use the output of that invocation
        state.visit(new class extends ThingTalk.Ast.NodeVisitor {
            visitInvocation(invocation) {
                const selector = invocation.selector;
                assert_1.default(selector instanceof thingtalk_1.Ast.DeviceSelector);
                const device = selector.kind;
                const domain = device.substring(device.lastIndexOf('.') + 1).toLowerCase();
                currentDomain = domain;
                // delete all slots for this domain (they'll be set again right after)
                for (const arg of invocation.schema.iterateArguments()) {
                    if (arg.name === currentDomain) {
                        // do not erase the "id" slot just because we have an action!
                        assert_1.default(arg.type.isEntity);
                        continue;
                    }
                    const slotKey = nameToSlot(domain, arg.name);
                    delete slotValues[slotKey];
                }
                for (const in_param of invocation.in_params) {
                    if (in_param.value.isUndefined || in_param.value.isVarRef)
                        continue;
                    const slotKey = nameToSlot(domain, in_param.name);
                    slotValues[slotKey] = in_param.value;
                }
                // do not recurse
                return false;
            }
            visitDialogueHistoryItem(item) {
                // recurse only if this item comes from the user and not the agent
                return item.confirm !== 'proposed';
            }
            visitDontCareBooleanExpression(expr) {
                const slotKey = nameToSlot(currentDomain, expr.name);
                slotValues[slotKey] = new ThingTalk.Ast.Value.Enum('dontcare');
                return false;
            }
            visitAtomBooleanExpression(expr) {
                if (expr.value.isUndefined || expr.value.isVarRef)
                    return false;
                const slotKey = nameToSlot(currentDomain, expr.name);
                if (expr.operator === 'in_array') // multiple values, pick the first one
                    slotValues[slotKey] = expr.value.value[0];
                else
                    slotValues[slotKey] = expr.value;
                return false;
            }
            visitNotBooleanExpression(expr) {
                // explicitly do not recurse into "not" operators
                return false;
            }
            visitOrBooleanExpression(expr) {
                // explicitly do not recurse into "or" operators
                // (unless they are an "or" of one operand)
                return expr.operands.length === 1;
            }
        });
        // remove train-name which is not in multiwoz
        delete slotValues['train-name'];
        // resolve entities and map Ast.Value to a string we can compare for equality
        const slotStrings = {};
        for (const key in slotValues)
            slotStrings[key] = await this._valueToSlot(slotValues[key]);
        return slotStrings;
    }
    async _checkTurn(id, turn, turnIndex) {
        let context, contextCode, contextEntities;
        if (turnIndex > 0) {
            if (turn.intermediate_context) {
                context = await ThingTalkUtils.parse(turn.intermediate_context, this._options);
                assert_1.default(context instanceof thingtalk_1.Ast.DialogueState);
            }
            else {
                context = await ThingTalkUtils.parse(turn.context, this._options);
                assert_1.default(context instanceof thingtalk_1.Ast.DialogueState);
                // apply the agent prediction to the context to get the state of the dialogue before
                // the user speaks
                const agentPrediction = await ThingTalkUtils.parse(turn.agent_target, this._options);
                assert_1.default(agentPrediction instanceof thingtalk_1.Ast.DialogueState);
                context = ThingTalkUtils.computeNewState(context, agentPrediction, 'agent');
            }
            const userContext = ThingTalkUtils.prepareContextForPrediction(context, 'user');
            [contextCode, contextEntities] = ThingTalkUtils.serializeNormalized(userContext);
        }
        else {
            context = null;
            contextCode = ['null'];
            contextEntities = {};
        }
        const { tokens, entities } = await this._preprocess(turn.user, contextEntities);
        const goldUserTarget = await ThingTalkUtils.parse(turn.user_target, this._options);
        assert_1.default(goldUserTarget instanceof thingtalk_1.Ast.DialogueState);
        const goldUserState = ThingTalkUtils.computeNewState(context, goldUserTarget, 'user');
        const goldSlots = await this._extractSlots(goldUserState);
        const targetCode = ThingTalkUtils.serializePrediction(goldUserTarget, tokens, entities, {
            locale: this._locale,
        }).join(' ');
        const parsed = await this._parser.sendUtterance(tokens.join(' '), contextCode, contextEntities, {
            tokenized: true,
            skip_typechecking: true,
            example_id: id + '/' + turnIndex,
        });
        const predictions = parsed.candidates
            .filter((beam) => beam.score !== 'Infinity') // ignore exact matches
            .map((beam) => beam.code);
        if (predictions.length === 0) {
            if (this._debug)
                console.log(`${id}:${turnIndex}\twrong_syntax\t${contextCode.join(' ')}\t${turn.user}\tfailed\t${targetCode}`);
            return 'wrong_syntax';
        }
        const choice = predictions[0];
        // first check if the program parses and typechecks (no hope otherwise)
        let predictedUserTarget;
        try {
            predictedUserTarget = await ThingTalkUtils.parsePrediction(choice, entities, this._options, true);
            assert_1.default(predictedUserTarget instanceof thingtalk_1.Ast.DialogueState);
        }
        catch (e) {
            if (this._debug)
                console.log(`${id}:${turnIndex}\twrong_syntax\t${contextCode.join(' ')}\t${turn.user}\t${choice.join(' ')}\t${targetCode}`);
            return 'wrong_syntax';
        }
        const predictedUserState = ThingTalkUtils.computeNewState(context, predictedUserTarget, 'user');
        let predictedSlots;
        try {
            predictedSlots = await this._extractSlots(predictedUserState);
        }
        catch (e) {
            console.error(predictedUserTarget.prettyprint());
            throw e;
        }
        // do some light syntactic normalization
        const choiceString = Array.from(eval_utils_1.stripOutTypeAnnotations(choice)).join(' ');
        // do the actual normalization, using the full ThingTalk algorithm
        // we pass "ignoreSentence: true", which means strings are tokenized and then put in the
        // program regardless of what the sentence contains (because the neural network might
        // get creative in copying, and we don't want to crash here)
        const normalized = ThingTalkUtils.serializePrediction(predictedUserTarget, tokens, entities, {
            locale: this._locale,
            ignoreSentence: true
        }).join(' ');
        // check that by normalizing we did not accidentally mark wrong a program that
        // was correct before
        if (choiceString === targetCode && normalized !== targetCode) {
            console.error();
            console.error('NORMALIZATION ERROR');
            console.error(targetCode);
            console.error(normalized);
            console.error(choice.join(' '));
            console.error(choiceString);
            throw new Error('Normalization Error');
        }
        if (normalized === targetCode) {
            if (this._debug)
                console.log(`${id}:${turnIndex}\tok\t${contextCode.join(' ')}\t${turn.user}\t${normalized}\t${targetCode}`);
            if (!deep_equal_1.default(goldSlots, predictedSlots, { strict: true })) {
                console.error(goldSlots, predictedSlots);
                throw new Error(`Program matches but slots do not`);
            }
            return 'ok';
        }
        else if (deep_equal_1.default(goldSlots, predictedSlots, { strict: true })) {
            if (this._debug)
                console.log(`${id}:${turnIndex}\tok_slot\t${contextCode.join(' ')}\t${turn.user}\t${normalized}\t${targetCode}`);
            return 'ok_slot';
        }
        else {
            if (this._debug)
                console.log(`${id}:${turnIndex}\tok_syntax\t${contextCode.join(' ')}\t${turn.user}\t${normalized}\t${targetCode}`);
            return 'ok_syntax';
        }
    }
    async _evaluate(dialogue) {
        let prefix_full = 0;
        let correct_full = 0;
        let prefix_slot = 0;
        let correct_slot = 0;
        let failed_full = false, failed_slot = false;
        for (let i = 0; i < dialogue.length; i++) {
            const turn = dialogue[i];
            let ok;
            try {
                ok = await this._checkTurn(dialogue.id, turn, i);
            }
            catch (e) {
                console.error(dialogue.id, turn);
                throw e;
            }
            if (ok === 'ok') {
                correct_full += 1;
                correct_slot += 1;
                if (!failed_slot)
                    prefix_slot += 1;
                if (!failed_full)
                    prefix_full += 1;
            }
            else if (ok === 'ok_slot') {
                correct_slot += 1;
                if (!failed_slot)
                    prefix_slot += 1;
                failed_full = true;
            }
            else {
                failed_full = true;
                failed_slot = true;
            }
        }
        const ret = {
            turns: dialogue.length,
            ok: Number(correct_full === dialogue.length),
            ok_slot: Number(correct_slot === dialogue.length),
            ok_initial: Number(prefix_full >= 1),
            ok_initial_slot: Number(prefix_slot >= 1),
            ok_partial: Number(correct_full),
            ok_partial_slot: Number(correct_slot),
            ok_prefix: Number(prefix_full),
            ok_prefix_slot: Number(prefix_slot),
            ok_progress: Number(prefix_full),
            ok_progress_slot: Number(prefix_slot),
        };
        if (this._debug) {
            let message = String(dialogue.id);
            for (const key in ret) {
                if (key === 'turns')
                    continue;
                message += '\t' + ret[key];
            }
            console.log(message);
        }
        return ret;
    }
    async _flushMinibatch() {
        for (const res of await Promise.all(this._minibatch))
            this.push(res);
        this._minibatch = [];
    }
    async _pushDialogue(dialog) {
        this._minibatch.push(this._evaluate(dialog));
        if (this._minibatch.length >= MINIBATCH_SIZE)
            await this._flushMinibatch();
    }
    _transform(dialog, encoding, callback) {
        this._pushDialogue(dialog).then(() => callback(null), (err) => callback(err));
    }
    _flush(callback) {
        this._flushMinibatch().then(() => callback(null), (err) => callback(err));
    }
}
exports.DialogueEvaluatorStream = DialogueEvaluatorStream;
const KEYS = [
    'ok', 'ok_slot', 'ok_initial', 'ok_initial_slot',
    'ok_partial', 'ok_partial_slot', 'ok_prefix', 'ok_prefix_slot',
    'ok_progress', 'ok_progress_slot'
];
exports.KEYS = KEYS;
const BY_TURN_KEYS = ['ok_partial', 'ok_partial_slot', 'ok_prefix', 'ok_prefix_slot'];
const BY_DIALOGUE_KEYS = ['ok', 'ok_slot', 'ok_initial', 'ok_initial_slot', 'ok_progress', 'ok_progress_slot'];
class CollectDialogueStatistics extends stream_1.default.Writable {
    constructor() {
        super({ objectMode: true });
        this._buffer = {
            total: 0,
            turns: 0,
        };
        for (const key of KEYS)
            this._buffer[key] = 0;
    }
    _write(sample, encoding, callback) {
        this._buffer.total++;
        this._buffer.turns += sample.turns;
        for (const key of KEYS)
            this._buffer[key] += sample[key];
        callback();
    }
    _final(callback) {
        // convert to percentages
        for (const key of BY_DIALOGUE_KEYS)
            this._buffer[key] /= this._buffer.total;
        for (const key of BY_TURN_KEYS)
            this._buffer[key] /= this._buffer.turns;
        callback();
    }
    read() {
        return new Promise((resolve, reject) => {
            this.on('finish', () => resolve(this._buffer));
            this.on('error', reject);
        });
    }
}
exports.CollectDialogueStatistics = CollectDialogueStatistics;
//# sourceMappingURL=dialogue_evaluator.js.map