"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());
    });
};
var __generator = (this && this.__generator) || function (thisArg, body) {
    var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g;
    return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g;
    function verb(n) { return function (v) { return step([n, v]); }; }
    function step(op) {
        if (f) throw new TypeError("Generator is already executing.");
        while (_) try {
            if (f = 1, y && (t = y[op[0] & 2 ? "return" : op[0] ? "throw" : "next"]) && !(t = t.call(y, op[1])).done) return t;
            if (y = 0, t) op = [0, t.value];
            switch (op[0]) {
                case 0: case 1: t = op; break;
                case 4: _.label++; return { value: op[1], done: false };
                case 5: _.label++; y = op[1]; op = [0]; continue;
                case 7: op = _.ops.pop(); _.trys.pop(); continue;
                default:
                    if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; }
                    if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; }
                    if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; }
                    if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; }
                    if (t[2]) _.ops.pop();
                    _.trys.pop(); continue;
            }
            op = body.call(thisArg, _);
        } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; }
        if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true };
    }
};
Object.defineProperty(exports, "__esModule", { value: true });
var tfc = require("@tensorflow/tfjs-core");
var state_1 = require("../backend/state");
var base_callbacks_1 = require("../base_callbacks");
var errors_1 = require("../errors");
var logs_1 = require("../logs");
var generic_utils_1 = require("../utils/generic_utils");
var DEFAULT_VALIDATION_BATCH_SIZE = 32;
function standardizeDataIteratorOutput(model, iteratorOut) {
    tfc.util.assert(Array.isArray(iteratorOut) && iteratorOut.length === 2, 'Dataset iterator for fitDataset() is expected to generate ' +
        'an Array of length 2: `[xs, ys]`, but instead generates ' +
        iteratorOut);
    var tuple = iteratorOut;
    var xs = tuple[0];
    var ys = tuple[1];
    var standardizedXs = [];
    var standardizedYs = [];
    var xsBatchSize;
    var ysBatchSize;
    if (xs instanceof tfc.Tensor) {
        tfc.util.assert(model.inputs.length === 1, "Model has multiple " + model.inputs.length + " inputs, hence it " +
            "expects the input dataset to generate a dictionary of tensors " +
            ("(with keys " + JSON.stringify(model.inputNames) + "), ") +
            "but received a single tensor.");
        xsBatchSize = xs.shape[0];
        standardizedXs.push(xs);
    }
    else {
        xs = xs;
        for (var _i = 0, _a = model.inputNames; _i < _a.length; _i++) {
            var inputName = _a[_i];
            if (xs[inputName] == null) {
                throw new errors_1.ValueError("The feature data generated by the dataset " +
                    ("lacks the required input key '" + inputName + "'."));
            }
            if (xsBatchSize == null) {
                xsBatchSize = xs[inputName].shape[0];
            }
            else {
                tfc.util.assert(xs[inputName].shape[0] === xsBatchSize, "Mismatch in batch size between inputs x tensors, " +
                    ("(" + xs[inputName].shape[0] + " vs. " + xsBatchSize + ")"));
            }
            standardizedXs.push(xs[inputName]);
        }
    }
    if (ys instanceof tfc.Tensor) {
        tfc.util.assert(model.outputs.length === 1, "Model has multiple " + model.outputs.length + " outputs, hence it " +
            "expects the output dataset to generate a dictionary of tensors " +
            ("(with keys " + JSON.stringify(model.outputNames) + "), ") +
            "but received a single tensor.");
        ysBatchSize = ys.shape[0];
        standardizedYs.push(ys);
    }
    else {
        ys = ys;
        for (var _b = 0, _c = model.outputNames; _b < _c.length; _b++) {
            var outputName = _c[_b];
            if (ys[outputName] == null) {
                throw new errors_1.ValueError("The feature data generated by the dataset " +
                    ("lacks the required output key '" + outputName + "'."));
            }
            if (ysBatchSize == null) {
                ysBatchSize = ys[outputName].shape[0];
            }
            else {
                tfc.util.assert(ys[outputName].shape[0] === ysBatchSize, "Mismatch in batch size between output y tensors, " +
                    ("(" + ys[outputName].shape[0] + " vs. " + ysBatchSize + ")"));
            }
            standardizedYs.push(ys[outputName]);
        }
    }
    tfc.util.assert(xsBatchSize === ysBatchSize, "Mismatch in batch size between x and y tensors (" + xsBatchSize + " vs. " +
        (ysBatchSize + ")"));
    return standardizedXs.concat(standardizedYs);
}
function standardizeTensorValidationData(data) {
    if (data.length === 3) {
        throw new errors_1.NotImplementedError('Validation with sample weights is not implemented yet.');
    }
    return { xs: data[0], ys: data[1] };
}
function fitDataset(model, dataset, args) {
    return __awaiter(this, void 0, void 0, function () {
        var hasBatchesPerEpoch, doValidation, valXs, valYs, validationData, trainFunction, outLabels, callbackMetrics, callbacks, verbose, _a, callbackList, history_1, epoch, dataIterator, epochLogs, stepsDone, batchIndex, iteratorOut, xsAndYs, batchLogs, outs, i, label, out, valOuts, _b, i;
        return __generator(this, function (_c) {
            switch (_c.label) {
                case 0:
                    hasBatchesPerEpoch = args.batchesPerEpoch != null;
                    tfc.util.assert(model.optimizer != null, 'You must compile a model before training/testing. Use ' +
                        'Model.compile(modelCompileConfig).');
                    tfc.util.assert(args != null, "For fitDataset(), the 2nd argument (config) is required, " +
                        "but it is not provided in this call.");
                    tfc.util.assert(args.epochs != null && args.epochs > 0 && Number.isInteger(args.epochs), "For fitDataset(), config.epochs is expected to be a positive " +
                        ("integer, but got " + args.epochs));
                    tfc.util.assert(!hasBatchesPerEpoch ||
                        (args.batchesPerEpoch > 0 && Number.isInteger(args.batchesPerEpoch)), "For fitDataset(), config.batchesPerEpoch is expected to be a " +
                        ("positive integer if specified, but got " + args.batchesPerEpoch));
                    tfc.util.assert(args['validationSplit'] == null, '`validationSplit` is not supported by `fitDataset()`. ' +
                        'Use validationData instead.');
                    if (model.isTraining) {
                        throw new Error('Cannot start training because another fit() call is ongoing.');
                    }
                    model.isTraining = true;
                    _c.label = 1;
                case 1:
                    _c.trys.push([1, , 22, 23]);
                    doValidation = args.validationData != null;
                    valXs = void 0;
                    valYs = void 0;
                    if (doValidation) {
                        if (isDatasetObject(args.validationData)) {
                            tfc.util.assert(args.validationBatches == null ||
                                (args.validationBatches > 0 &&
                                    Number.isInteger(args.validationBatches)), "For fitDataset() with dataset-based validation, " +
                                "config.validationBatches is expected not to be provided, " +
                                "or to be a positive integer, " +
                                ("but got " + args.validationBatches));
                        }
                        else {
                            validationData = standardizeTensorValidationData(args.validationData);
                            valXs = validationData.xs;
                            valYs = validationData.ys;
                        }
                    }
                    trainFunction = model.makeTrainFunction();
                    outLabels = model.getDedupedMetricsNames();
                    callbackMetrics = void 0;
                    if (doValidation) {
                        callbackMetrics =
                            outLabels.slice().concat(outLabels.map(function (n) { return 'val_' + n; }));
                    }
                    else {
                        callbackMetrics = outLabels.slice();
                    }
                    callbacks = base_callbacks_1.standardizeCallbacks(args.callbacks);
                    verbose = args.verbose == null ? 1 : args.verbose;
                    _a = base_callbacks_1.configureCallbacks(callbacks, args.yieldEvery, verbose, args.epochs, null, null, getStepsPerEpoch(dataset, args), null, doValidation, callbackMetrics), callbackList = _a.callbackList, history_1 = _a.history;
                    callbackList.setModel(model);
                    model.history = history_1;
                    return [4, callbackList.onTrainBegin()];
                case 2:
                    _c.sent();
                    model.stopTraining_ = false;
                    epoch = args.initialEpoch == null ? 0 : args.initialEpoch;
                    return [4, dataset.iterator()];
                case 3:
                    dataIterator = _c.sent();
                    _c.label = 4;
                case 4:
                    if (!(epoch < args.epochs)) return [3, 19];
                    epochLogs = {};
                    return [4, callbackList.onEpochBegin(epoch)];
                case 5:
                    _c.sent();
                    stepsDone = 0;
                    batchIndex = 0;
                    if (!!hasBatchesPerEpoch) return [3, 7];
                    return [4, dataset.iterator()];
                case 6:
                    dataIterator = _c.sent();
                    _c.label = 7;
                case 7:
                    if (!(hasBatchesPerEpoch ? stepsDone < args.batchesPerEpoch : true)) return [3, 17];
                    return [4, dataIterator.next()];
                case 8:
                    iteratorOut = _c.sent();
                    if (hasBatchesPerEpoch && iteratorOut.done) {
                        console.warn('You provided `batchesPerEpoch` as ' +
                            (args.batchesPerEpoch + ", ") +
                            'but your dataset iterator ran out of data after ' +
                            (stepsDone + " batches; ") +
                            'interrupting training. Make sure that your ' +
                            'dataset can generate at least `batchesPerEpoch * epochs` ' +
                            'batches (in this case, ' +
                            (args.batchesPerEpoch * args.epochs + " batches). ") +
                            'You may need to use the repeat() function when building ' +
                            'your dataset.');
                        return [3, 17];
                    }
                    if (!(iteratorOut.value != null)) return [3, 11];
                    xsAndYs = standardizeDataIteratorOutput(model, iteratorOut.value);
                    batchLogs = {};
                    batchLogs['batch'] = batchIndex;
                    batchLogs['size'] = xsAndYs[0].shape[0];
                    return [4, callbackList.onBatchBegin(batchIndex, batchLogs)];
                case 9:
                    _c.sent();
                    outs = trainFunction(xsAndYs);
                    tfc.dispose(xsAndYs);
                    for (i = 0; i < outLabels.length; ++i) {
                        label = outLabels[i];
                        out = outs[i];
                        batchLogs[label] = out;
                        tfc.keep(out);
                    }
                    return [4, callbackList.onBatchEnd(batchIndex, batchLogs)];
                case 10:
                    _c.sent();
                    logs_1.disposeTensorsInLogs(batchLogs);
                    batchIndex++;
                    stepsDone++;
                    _c.label = 11;
                case 11:
                    if (!(hasBatchesPerEpoch ? stepsDone >= args.batchesPerEpoch :
                        iteratorOut.done)) return [3, 16];
                    if (!doValidation) return [3, 15];
                    valOuts = void 0;
                    if (!isDatasetObject(args.validationData)) return [3, 13];
                    _b = generic_utils_1.toList;
                    return [4, model.evaluateDataset(args.validationData, { batches: args.validationBatches })];
                case 12:
                    valOuts = _b.apply(void 0, [_c.sent()]);
                    return [3, 14];
                case 13:
                    valOuts = generic_utils_1.toList(model.evaluate(valXs, valYs, {
                        batchSize: args.validationBatchSize == null ?
                            DEFAULT_VALIDATION_BATCH_SIZE :
                            args.validationBatchSize,
                        verbose: 0
                    }));
                    _c.label = 14;
                case 14:
                    for (i = 0; i < model.metricsNames.length; ++i) {
                        epochLogs["val_" + model.metricsNames[i]] = valOuts[i];
                    }
                    _c.label = 15;
                case 15: return [3, 17];
                case 16:
                    if (model.stopTraining_) {
                        return [3, 17];
                    }
                    return [3, 7];
                case 17: return [4, callbackList.onEpochEnd(epoch, epochLogs)];
                case 18:
                    _c.sent();
                    epoch++;
                    if (model.stopTraining_) {
                        return [3, 19];
                    }
                    return [3, 4];
                case 19: return [4, callbackList.onTrainEnd()];
                case 20:
                    _c.sent();
                    return [4, model.history.syncData()];
                case 21:
                    _c.sent();
                    return [2, model.history];
                case 22:
                    model.isTraining = false;
                    return [7];
                case 23: return [2];
            }
        });
    });
}
exports.fitDataset = fitDataset;
function getStepsPerEpoch(dataset, args) {
    var stepsPerEpoch = null;
    if (args.batchesPerEpoch != null) {
        stepsPerEpoch = args.batchesPerEpoch;
    }
    else if (Number.isFinite(dataset.size)) {
        stepsPerEpoch = dataset.size;
    }
    return stepsPerEpoch;
}
function isDatasetObject(dataset) {
    return (typeof dataset.iterator === 'function');
}
function isLazyIteratorObject(iterator) {
    return (typeof iterator.next === 'function');
}
function evaluateDataset(model, dataset, args) {
    return __awaiter(this, void 0, void 0, function () {
        var hasBatches, f, outs, dataIterator, _a, numExamples, batch, _loop_1, state_2, _loop_2, i;
        return __generator(this, function (_b) {
            switch (_b.label) {
                case 0:
                    args = args || {};
                    hasBatches = args.batches != null;
                    f = model.testFunction;
                    outs = [];
                    if (args.verbose > 0) {
                        throw new errors_1.NotImplementedError('Verbose mode is not implemented yet.');
                    }
                    tfc.util.assert(!hasBatches || (args.batches > 0 && Number.isInteger(args.batches)), 'Test loop expects `batches` to be a positive integer, but ' +
                        ("received " + JSON.stringify(args.batches)));
                    if (!isLazyIteratorObject(dataset)) return [3, 1];
                    _a = dataset;
                    return [3, 3];
                case 1: return [4, dataset.iterator()];
                case 2:
                    _a = _b.sent();
                    _b.label = 3;
                case 3:
                    dataIterator = _a;
                    numExamples = 0;
                    batch = 0;
                    _loop_1 = function () {
                        var iteratorOut, xsAndYs_1, batchOuts, i, batchSize_1, _loop_3, i;
                        return __generator(this, function (_a) {
                            switch (_a.label) {
                                case 0: return [4, dataIterator.next()];
                                case 1:
                                    iteratorOut = _a.sent();
                                    if (iteratorOut.value) {
                                        xsAndYs_1 = standardizeDataIteratorOutput(model, iteratorOut.value);
                                        batchOuts = tfc.tidy(function () { return f(xsAndYs_1); });
                                        tfc.dispose(xsAndYs_1);
                                        if (batch === 0) {
                                            for (i = 0; i < batchOuts.length; ++i) {
                                                outs.push(state_1.getScalar(0));
                                            }
                                        }
                                        batchSize_1 = xsAndYs_1[0].shape[0];
                                        _loop_3 = function (i) {
                                            var batchOut = batchOuts[i];
                                            var oldScalar = outs[i];
                                            outs[i] = tfc.tidy(function () { return tfc.add(outs[i], tfc.mul(state_1.getScalar(batchSize_1), batchOut)); });
                                            if (batch > 0) {
                                                tfc.dispose(oldScalar);
                                            }
                                        };
                                        for (i = 0; i < batchOuts.length; ++i) {
                                            _loop_3(i);
                                        }
                                        tfc.dispose(batchOuts);
                                        numExamples += batchSize_1;
                                        ++batch;
                                    }
                                    if (iteratorOut.done) {
                                        if (hasBatches) {
                                            console.warn('Your dataset iterator ran out of data during evaluateDataset(). ' +
                                                'Interrupting evalution. Make sure that your ' +
                                                'dataset can generate at least `batches` ' +
                                                ("batches (in this case, " + args.batches + " batches). ") +
                                                'You may need to use the repeat() function when building ' +
                                                'your dataset.');
                                        }
                                        return [2, "break"];
                                    }
                                    return [2];
                            }
                        });
                    };
                    _b.label = 4;
                case 4:
                    if (!(hasBatches ? batch < args.batches : true)) return [3, 6];
                    return [5, _loop_1()];
                case 5:
                    state_2 = _b.sent();
                    if (state_2 === "break")
                        return [3, 6];
                    return [3, 4];
                case 6:
                    _loop_2 = function (i) {
                        var oldScalar = outs[i];
                        outs[i] =
                            tfc.tidy(function () { return tfc.div(outs[i], state_1.getScalar(numExamples)); });
                        tfc.dispose(oldScalar);
                    };
                    for (i = 0; i < outs.length; ++i) {
                        _loop_2(i);
                    }
                    return [2, generic_utils_1.singletonOrArray(outs)];
            }
        });
    });
}
exports.evaluateDataset = evaluateDataset;
//# sourceMappingURL=training_dataset.js.map