"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.DatasetLoader = void 0;
/* eslint-disable @typescript-eslint/naming-convention */
var fs_1 = require("fs");
var js_yaml_1 = require("js-yaml");
var path_1 = require("path");
var array_utils_1 = require("../utils/array-utils");
var data_augmenter_1 = require("./data-augmenter");
var dataset_1 = require("./dataset");
var entities_parser_1 = require("./entities-parser");
var YAML_EXTENSION = '.yaml';
var YML_EXTENSION = '.yml';
var DatasetLoader = /** @class */ (function () {
    function DatasetLoader() {
    }
    DatasetLoader.load = function (path) {
        var stat = fs_1.lstatSync(path);
        if (stat.isFile()) {
            return this.loadFile(path);
        }
        if (stat.isDirectory()) {
            return this.loadDirectory(path);
        }
        throw new Error("path \"" + path + "\" must be a directory or file.");
    };
    DatasetLoader.loadDirectory = function (path) {
        var _this = this;
        var files = fs_1.readdirSync(path);
        var datasets = files.map(function (filePath) { return _this.loadFile(path_1.join(path, filePath)); });
        var classes = array_utils_1.unique(datasets.reduce(function (classes, dataset) { return classes.concat(dataset.classes); }, []));
        var entities = array_utils_1.unique(datasets.reduce(function (entities, dataset) { return entities.concat(dataset.entities); }, []));
        var samples = datasets.reduce(function (samples, dataset) { return samples.concat(dataset.samples); }, []);
        return new dataset_1.Dataset(classes, entities, samples);
    };
    DatasetLoader.loadFile = function (path) {
        if (!this.ALLOWED_FILE_EXTENSIONS.includes(path_1.extname(path))) {
            throw new Error("File '" + path + "' must be a " + this.ALLOWED_FILE_EXTENSIONS.join(',') + ".");
        }
        var content = js_yaml_1.load(fs_1.readFileSync(path));
        if (!(this.SAMPLES_FIELD in content)) {
            throw new Error('File must contain "samples" field');
        }
        var className = this.CLASS_FIELD in content ? content[this.CLASS_FIELD] : '';
        var classes = this.CLASS_FIELD in content ? [className] : [];
        var entities = this.ENTITIES_FIELD in content ? array_utils_1.unique(content[this.ENTITIES_FIELD]) : [];
        var sentences = content[this.SAMPLES_FIELD];
        if (this.DATA_AUGMENTATION_FIELD in content) {
            var augmenter = new data_augmenter_1.DataAugmenter(content[this.DATA_AUGMENTATION_FIELD], entities);
            sentences = augmenter.augment(sentences);
        }
        var samples = sentences.map(function (s) {
            return { text: s, class: className, entities: [] };
        });
        if (this.ENTITIES_FIELD in content) {
            samples = entities_parser_1.EntitiesParser.parse(samples, entities);
        }
        return new dataset_1.Dataset(classes, entities, samples);
    };
    DatasetLoader.ALLOWED_FILE_EXTENSIONS = [YAML_EXTENSION, YML_EXTENSION];
    DatasetLoader.CLASS_FIELD = 'class';
    DatasetLoader.ENTITIES_FIELD = 'entities';
    DatasetLoader.DATA_AUGMENTATION_FIELD = 'data-augmentation';
    DatasetLoader.SAMPLES_FIELD = 'samples';
    return DatasetLoader;
}());
exports.DatasetLoader = DatasetLoader;
//# sourceMappingURL=loader.js.map