"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : new P(function (resolve) { resolve(result.value); }).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
Object.defineProperty(exports, "__esModule", { value: true });
const chance_1 = require("chance");
const logger = console;
exports.VALID_DISTRIBUTIONS = ['regular', 'even'];
exports.config = {
    defaultDistribution: 'regular'
};
const chatito = require('../parser/chatito');
const chance = new chance_1.Chance();
const getEntityKey = (token) => (token.variation ? `${token.value}#${token.variation}` : token.value);
const chatitoFormatPostProcess = (data) => {
    const arr = data.reduce((accumulator, next, i, arrShadow) => {
        if (accumulator.length) {
            const lastWord = accumulator[accumulator.length - 1];
            if (lastWord.type === next.type && lastWord.type === 'Text') {
                accumulator[accumulator.length - 1] = {
                    type: lastWord.type,
                    value: (lastWord.value + next.value).replace(/\s+/g, ' ')
                };
            }
            else {
                accumulator.push(next);
            }
        }
        else if (next.value.trim()) {
            accumulator.push(next);
        }
        if (i === arrShadow.length - 1) {
            if (accumulator.length) {
                if (!accumulator[accumulator.length - 1].value.trim()) {
                    accumulator.pop();
                }
                accumulator[accumulator.length - 1] = Object.assign({}, accumulator[accumulator.length - 1], {
                    value: accumulator[accumulator.length - 1].value.replace(/\s+$/g, '')
                });
            }
        }
        return accumulator;
    }, []);
    if (arr.length) {
        arr[0] = Object.assign({}, arr[0], {
            value: arr[0].value.replace(/^\s+/, '')
        });
    }
    if (!arr.length) {
        throw new Error(`Some sentence generated an empty string. Can't map empty to an intent.`);
    }
    return arr;
};
const calcSentencesProbabilities = (isPercentageProbability, isEvenDistribution, definedSentenceProbabilities, sumOfTotalProbabilitiesDefined, maxCounts) => {
    let sentencesWithNullProbabilityCount = 0;
    let totalMaxCountsToShareBetweenNullProbs = 0;
    definedSentenceProbabilities.forEach((prob, i) => {
        if (prob === null) {
            sentencesWithNullProbabilityCount += 1;
            totalMaxCountsToShareBetweenNullProbs += maxCounts[i];
        }
    });
    let probabilities;
    if (isPercentageProbability) {
        probabilities = definedSentenceProbabilities.map((p, i) => {
            if (p !== null) {
                return p;
            }
            if (isEvenDistribution) {
                return (100 - sumOfTotalProbabilitiesDefined) / sentencesWithNullProbabilityCount;
            }
            return (((maxCounts[i] * 100) / totalMaxCountsToShareBetweenNullProbs) * (100 - sumOfTotalProbabilitiesDefined)) / 100;
        });
    }
    else {
        probabilities = definedSentenceProbabilities.map((p, i) => {
            if (p !== null) {
                return isEvenDistribution ? p : maxCounts[i] * p;
            }
            if (isEvenDistribution) {
                return 1;
            }
            return maxCounts[i];
        });
    }
    return probabilities;
};
exports.getVariationsFromEntity = (ed, entities, optional, cache) => __awaiter(this, void 0, void 0, function* () {
    const variationKey = ed.variation ? `#${ed.variation}` : '';
    const cacheKey = `${ed.type}-${ed.key}${variationKey}`;
    let cacheStats = cache.get(cacheKey);
    if (!cacheStats) {
        const counts = [];
        const maxCounts = ed.inner.map(s => s.cardinality);
        let probabilityTypeDefined = null;
        const definedSentenceProbabilities = [];
        let isEvenDistribution = exports.config.defaultDistribution === 'even';
        if (ed.args && ed.args.distribution) {
            isEvenDistribution = ed.args.distribution === 'even';
        }
        let sumOfTotalProbabilitiesDefined = 0;
        for (const c of ed.inner) {
            counts.push(new Map());
            if (c.probability === null) {
                definedSentenceProbabilities.push(null);
            }
            else {
                const p = c.probability || '';
                const isPercent = p.slice(-1) === '%';
                const setenceProbabilityType = isPercent ? '%' : 'w';
                if (probabilityTypeDefined === null) {
                    probabilityTypeDefined = setenceProbabilityType;
                }
                else if (setenceProbabilityType !== probabilityTypeDefined) {
                    throw new Error(`All probability definitions for "${cacheKey}" must be of the same type.`);
                }
                const prob = parseFloat(isPercent ? p.slice(0, -1) : p);
                if (isPercent) {
                    if (prob <= 0 || prob > 100) {
                        throw new Error(`Probability "${p}" must be greater than 0 up to 100. At ${cacheKey}`);
                    }
                }
                else if (setenceProbabilityType === 'w') {
                    if (prob <= 0) {
                        throw new Error(`Probability weight "${p}" must be greater than 0. At ${cacheKey}`);
                    }
                }
                sumOfTotalProbabilitiesDefined += prob;
                definedSentenceProbabilities.push(prob);
            }
        }
        if (probabilityTypeDefined === '%' && sumOfTotalProbabilitiesDefined && sumOfTotalProbabilitiesDefined > 100) {
            throw new Error(`The sum of sentence probabilities (${sumOfTotalProbabilitiesDefined}) for an entity can't be higher than 100%. At ${cacheKey}`);
        }
        const isPercentageProbability = probabilityTypeDefined === '%';
        const probabilities = calcSentencesProbabilities(isPercentageProbability, isEvenDistribution, definedSentenceProbabilities, sumOfTotalProbabilitiesDefined, maxCounts);
        const currentEntityCache = { counts, maxCounts, probabilities };
        cache.set(cacheKey, currentEntityCache);
        cacheStats = cache.get(cacheKey);
    }
    const optionalProb = 100 / (cacheStats.probabilities.length + 1);
    const sentenceIndex = chance.weighted(Array.from(cacheStats.probabilities.keys()), cacheStats.probabilities);
    if (optional && chance.bool({ likelihood: optionalProb })) {
        return [];
    }
    const sentence = ed.inner[sentenceIndex].sentence;
    let accumulator = [];
    const isSlotDefSentenceWithOnlyOneAlias = ed.type === 'SlotDefinition' && sentence.length === 1 && sentence[0].type === 'Alias';
    for (const t of sentence) {
        const slotsInSentenceKeys = new Set([]);
        if (t.type === 'Slot' || t.type === 'Alias') {
            const def = entities[t.type];
            const innerEntityKey = getEntityKey(t);
            const currentCache = slotsInSentenceKeys.has(innerEntityKey) ? cacheStats.counts[sentenceIndex] : new Map();
            slotsInSentenceKeys.add(innerEntityKey);
            const sentenceVariation = yield exports.getVariationsFromEntity(def[innerEntityKey], entities, !!t.opt, currentCache);
            if (sentenceVariation.length) {
                const returnSentenceTokens = chatitoFormatPostProcess(sentenceVariation);
                for (const returnToken of returnSentenceTokens) {
                    const ettArgs = def[innerEntityKey].args;
                    if (isSlotDefSentenceWithOnlyOneAlias && ettArgs && ettArgs.synonym === 'true') {
                        returnToken.synonym = t.value;
                    }
                    if (t.type === 'Slot') {
                        if (def[innerEntityKey].args) {
                            returnToken.args = def[innerEntityKey].args;
                        }
                        returnToken.value = returnToken.value.trim();
                        returnToken.type = t.type;
                        returnToken.slot = t.value;
                    }
                    accumulator = accumulator.concat(returnToken);
                }
            }
        }
        else {
            accumulator = accumulator.concat(t);
        }
    }
    return accumulator;
});
exports.getExampleByNumber = (defs, entity, combinationNumber) => {
    let lookupNumber = combinationNumber;
    const sentence = entity.inner.find(s => {
        if (lookupNumber < s.cardinality) {
            return true;
        }
        lookupNumber -= s.cardinality;
        return false;
    });
    if (!sentence) {
        return [];
    }
    let prevCardinality = 1;
    let prevRemaining = 0;
    const isSlotDefSentenceWithOnlyOneAlias = entity.type === 'SlotDefinition' && sentence.sentence.length === 1 && sentence.sentence[0].type === 'Alias';
    const resultTokens = sentence.sentence.reduce((example, token) => {
        if (token.type === 'Text') {
            return example.concat([token]);
        }
        if (token.type === 'Slot' || token.type === 'Alias') {
            let cardinality = token.opt ? 1 : 0;
            const innerEntity = token.type === 'Alias' ? defs.Alias : defs.Slot;
            const entityKey = getEntityKey(token);
            cardinality += innerEntity[entityKey].cardinality;
            lookupNumber = (lookupNumber - prevRemaining) / prevCardinality;
            prevRemaining = lookupNumber % cardinality;
            prevCardinality = cardinality;
            if (prevRemaining === 0 && token.opt) {
                return example;
            }
            const innerNumber = token.opt ? prevRemaining - 1 : prevRemaining;
            let tokens = exports.getExampleByNumber(defs, innerEntity[entityKey], innerNumber);
            tokens = chatitoFormatPostProcess(tokens).map(t => {
                const ettArgs = innerEntity[entityKey].args;
                if (isSlotDefSentenceWithOnlyOneAlias && ettArgs && ettArgs.synonym === 'true') {
                    t.synonym = token.value;
                }
                if (token.type === 'Slot') {
                    if (innerEntity[entityKey].args) {
                        t.args = innerEntity[entityKey].args;
                    }
                    t.value = t.value.trim();
                    t.type = token.type;
                    t.slot = token.value;
                }
                return t;
            });
            return example.concat(tokens);
        }
        throw Error(`Unknown token type: ${token.type}`);
    }, []);
    return chatitoFormatPostProcess(resultTokens);
};
function* allExamplesGenerator(defs, entity) {
    for (let i = 0; i < entity.cardinality; i++) {
        yield exports.getExampleByNumber(defs, entity, i);
    }
}
exports.allExamplesGenerator = allExamplesGenerator;
const getCardinality = (defs, sentence) => {
    return sentence.reduce((acc, token) => {
        if (token.type === 'Text') {
            return acc;
        }
        const entity = token.type === 'Alias' ? defs.Alias : defs.Slot;
        const entityKey = getEntityKey(token);
        let tokenCardinality = entity[entityKey].cardinality;
        if (token.opt) {
            tokenCardinality += 1;
        }
        return acc * tokenCardinality;
    }, 1);
};
const calcCardinality = (defs, entity) => {
    entity.inner.forEach(sentence => {
        const cardinality = getCardinality(defs, sentence.sentence);
        sentence.cardinality = cardinality;
    });
    entity.cardinality = entity.inner.reduce((acc, sentence) => acc + sentence.cardinality, 0);
};
const getRefKey = (item) => {
    const type = item.type.replace('Definition', '');
    const key = 'key' in item ? item.key : getEntityKey(item);
    switch (type) {
        case 'Intent':
            return `%[${key}]`;
        case 'Alias':
            return `~[${key}]`;
        case 'Slot':
            return `@[${key}]`;
        default:
            return `(${key})`;
    }
};
const hasTokenWithoutCardinality = (defs, entity, refs) => {
    const parentKey = getRefKey(entity);
    return entity.inner.some(sentence => sentence.sentence.some(token => {
        if (token.type === 'Text') {
            return false;
        }
        const entityKey = getEntityKey(token);
        const refKey = getRefKey(token);
        if (refKey in refs) {
            refs[refKey].add(parentKey);
        }
        else {
            refs[refKey] = new Set([parentKey]);
        }
        if (!defs[token.type][entityKey]) {
            throw new Error(`${token.type} not defined: ${entityKey}`);
        }
        return defs[token.type][entityKey].cardinality === undefined;
    }));
};
const checkLoopIn = (path, refs) => {
    const last = path[path.length - 1];
    if (refs[last]) {
        for (const parent of refs[last]) {
            if (parent === path[0]) {
                const loop = path.concat([parent]).reverse();
                throw new Error(`You have a circular nesting: ${loop.join(' -> ')}. Infinite loop prevented.`);
            }
            else {
                checkLoopIn(path.concat([parent]), refs);
            }
        }
    }
};
const checkLoop = (refs) => {
    for (const key of Object.keys(refs)) {
        const path = [key];
        checkLoopIn(path, refs);
    }
};
const findNestedSlots = (path, refs) => {
    const last = path[path.length - 1];
    if (refs[last]) {
        for (const parent of refs[last]) {
            const firstIndex = path.findIndex(item => item.startsWith('@'));
            if (firstIndex !== -1 && parent.startsWith('@')) {
                const slotsPath = path
                    .slice(firstIndex)
                    .reverse()
                    .join(' -> ');
                throw new Error(`You have nested slots: ${parent} -> ${slotsPath}. A slot can't reference other slot.`);
            }
            else {
                findNestedSlots(path.concat([parent]), refs);
            }
        }
    }
};
const checkNestedSlots = (refs) => {
    for (const key of Object.keys(refs)) {
        const path = [key];
        findNestedSlots(path, refs);
    }
};
const preCalcCardinality = (defs) => {
    const uncalced = {
        Intent: [],
        Alias: [],
        Slot: []
    };
    const refs = {};
    let totalUncalced = 0;
    let lastUncalced = -1;
    do {
        totalUncalced = 0;
        for (const type of Object.keys(uncalced)) {
            uncalced[type] = Object.keys(defs[type]).filter(key => defs[type][key].cardinality === undefined);
            uncalced[type].forEach(key => {
                if (!hasTokenWithoutCardinality(defs, defs[type][key], refs)) {
                    calcCardinality(defs, defs[type][key]);
                }
                else {
                    totalUncalced += 1;
                }
            });
        }
        if (lastUncalced === totalUncalced) {
            checkLoop(refs);
        }
        lastUncalced = totalUncalced;
    } while (totalUncalced > 0);
    checkNestedSlots(refs);
};
const addMissingAliases = (defs) => {
    const aliases = new Set();
    for (const entities of [defs.Alias, defs.Slot, defs.Intent]) {
        for (const key of Object.keys(entities)) {
            entities[key].inner.forEach(sentence => {
                sentence.sentence.forEach(token => {
                    if (token.type === 'Alias') {
                        aliases.add(token.value);
                    }
                });
            });
        }
    }
    for (const alias of aliases) {
        if (!defs.Alias[alias]) {
            defs.Alias[alias] = {
                inner: [{ sentence: [{ value: alias, type: 'Text' }], probability: null }],
                key: alias,
                type: 'AliasDefinition'
            };
        }
    }
};
exports.astFromString = (str) => chatito.parse(str);
exports.datasetFromString = (str, writterFn, importer, currentPath) => {
    const ast = exports.astFromString(str);
    return exports.datasetFromAST(ast, writterFn, importer, currentPath);
};
exports.getImports = (from, to, importer) => {
    const fileContent = importer(from, to);
    if (!fileContent || !fileContent.dsl) {
        throw new Error(`Failed importing ${to}`);
    }
    try {
        const importAst = exports.astFromString(fileContent.dsl);
        let outAst = [];
        importAst.forEach(ett => {
            if (ett.type === 'ImportFile' && ett.value) {
                outAst = [...outAst, ...exports.getImports(fileContent.filePath, ett.value, importer)];
            }
            else if (ett.type === 'AliasDefinition' || ett.type === 'SlotDefinition') {
                outAst = [...outAst, ett];
            }
        });
        return outAst;
    }
    catch (e) {
        throw new Error(`Failed importing ${to}. ${e.message} - ${JSON.stringify(e.location)}`);
    }
};
exports.definitionsFromAST = (initialAst, importHandler, currPath) => {
    const operatorDefinitions = { Intent: {}, Slot: {}, Alias: {} };
    if (!initialAst || !initialAst.length) {
        return;
    }
    const importer = importHandler ? importHandler : () => ({ filePath: '', dsl: '' });
    const currentPath = currPath ? currPath : '';
    let ast = [...initialAst];
    initialAst.forEach(od => {
        if (od.type === 'ImportFile' && od.value) {
            ast = [...ast, ...exports.getImports(currentPath, od.value, importer)];
        }
    });
    ast.forEach(od => {
        let entity;
        if (od.type === 'IntentDefinition') {
            entity = operatorDefinitions.Intent;
        }
        else if (od.type === 'SlotDefinition') {
            entity = operatorDefinitions.Slot;
        }
        else if (od.type === 'AliasDefinition') {
            entity = operatorDefinitions.Alias;
        }
        else {
            return;
        }
        const odKey = od.variation ? `${od.key}#${od.variation}` : od.key;
        if (entity[odKey]) {
            throw new Error(`Duplicate definition for ${od.type} '${odKey}'`);
        }
        entity[odKey] = od;
    });
    addMissingAliases(operatorDefinitions);
    preCalcCardinality(operatorDefinitions);
    return operatorDefinitions;
};
exports.datasetFromAST = (initialAst, writterFn, importHandler, currPath) => __awaiter(this, void 0, void 0, function* () {
    const operatorDefinitions = exports.definitionsFromAST(initialAst, importHandler, currPath);
    if (!operatorDefinitions) {
        return;
    }
    const intentKeys = Object.keys(operatorDefinitions.Intent);
    if (!intentKeys || !intentKeys.length) {
        return;
    }
    for (const intentKey of intentKeys) {
        const maxPossibleCombinations = operatorDefinitions.Intent[intentKey].cardinality;
        let maxIntentExamples = maxPossibleCombinations;
        const entityArgs = operatorDefinitions.Intent[intentKey].args;
        let trainingN = maxIntentExamples;
        let testingN = 0;
        let generatedTrainingExamplesCount = 0;
        let generatedTestingExamplesCount = 0;
        if (entityArgs) {
            if (entityArgs.training) {
                trainingN = parseInt(entityArgs.training, 10);
                if (trainingN < 1) {
                    throw new Error(`The 'training' argument for ${intentKey} must be higher than 0.`);
                }
                if (entityArgs.testing) {
                    testingN = parseInt(entityArgs.testing, 10);
                    if (testingN < 1) {
                        throw new Error(`The 'testing' argument for ${intentKey} must be higher than 0.`);
                    }
                }
            }
            let intentMax = trainingN + testingN;
            if (intentMax > maxIntentExamples) {
                logger.warn(`Can't generate ${intentMax} examples. ` +
                    `Using the maximum possible combinations: ${maxIntentExamples}. ` +
                    'NOTE: Using the maximum leads to overfitting.');
                intentMax = maxIntentExamples;
            }
            else if (intentMax < maxIntentExamples) {
                maxIntentExamples = intentMax;
            }
        }
        const maxEx = maxIntentExamples;
        const globalCache = new Map();
        const collitionsCache = {};
        if (maxIntentExamples >= maxPossibleCombinations) {
            for (const utterance of allExamplesGenerator(operatorDefinitions, operatorDefinitions.Intent[intentKey])) {
                const utteranceString = utterance.reduce((p, n) => p + n.value, '');
                if (!collitionsCache[utteranceString]) {
                    collitionsCache[utteranceString] = true;
                    const completedTraining = generatedTrainingExamplesCount >= trainingN;
                    const completedTesting = generatedTestingExamplesCount >= testingN;
                    let isTrainingExample = !completedTraining;
                    if (!completedTraining && !completedTesting) {
                        const trainingLeft = trainingN - generatedTrainingExamplesCount;
                        const testingLeft = testingN - generatedTestingExamplesCount;
                        isTrainingExample = Math.random() < trainingLeft / (trainingLeft + testingLeft);
                    }
                    writterFn(utterance, intentKey, isTrainingExample);
                    if (isTrainingExample) {
                        generatedTrainingExamplesCount++;
                    }
                    else {
                        generatedTestingExamplesCount++;
                    }
                }
            }
            continue;
        }
        let duplicatesCounter = 0;
        while (maxIntentExamples) {
            const intentSentence = yield exports.getVariationsFromEntity(operatorDefinitions.Intent[intentKey], operatorDefinitions, false, globalCache);
            const utterance = chatitoFormatPostProcess(intentSentence);
            const utteranceString = utterance.reduce((p, n) => p + n.value, '');
            if (!collitionsCache[utteranceString]) {
                collitionsCache[utteranceString] = true;
                const completedTraining = generatedTrainingExamplesCount >= trainingN;
                const completedTesting = generatedTestingExamplesCount >= testingN;
                let isTrainingExample = !completedTraining;
                if (!completedTraining && !completedTesting) {
                    isTrainingExample = Math.random() < 0.7;
                }
                writterFn(utterance, intentKey, isTrainingExample);
                maxIntentExamples--;
                if (isTrainingExample) {
                    generatedTrainingExamplesCount++;
                }
                else {
                    generatedTestingExamplesCount++;
                }
            }
            else {
                duplicatesCounter++;
                const smallDupesLimit = 10000;
                const maxDupes = maxPossibleCombinations * maxPossibleCombinations;
                const maxDupesLimit = Math.floor(maxDupes / 2);
                const isBigDataset = maxPossibleCombinations > smallDupesLimit;
                if ((isBigDataset && duplicatesCounter > maxDupesLimit) ||
                    (!isBigDataset && duplicatesCounter > maxDupes * maxPossibleCombinations)) {
                    let m = `Too many duplicates while generating dataset! Looks like we have probably reached `;
                    m += `the maximum ammount of possible unique generated examples. `;
                    m += `The generator has stopped at ${maxEx - maxIntentExamples} examples for intent ${intentKey}.`;
                    logger.warn(m);
                    maxIntentExamples = 0;
                }
            }
        }
    }
});
