import { nextFrame, tidy, image, scalar, layers, sequential, train, loadModel, model, oneHot, tensor1d, keep } from '@tensorflow/tfjs';

/*! *****************************************************************************
Copyright (c) Microsoft Corporation. All rights reserved.
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

THIS CODE IS PROVIDED ON AN *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
MERCHANTABLITY OR NON-INFRINGEMENT.

See the Apache Version 2.0 License for specific language governing permissions
and limitations under the License.
***************************************************************************** */

var __assign = Object.assign || function __assign(t) {
    for (var s, i = 1, n = arguments.length; i < n; i++) {
        s = arguments[i];
        for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p)) t[p] = s[p];
    }
    return t;
};

function __awaiter(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());
    });
}

function __generator(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 = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t;
            if (y = 0, t) op = [op[0] & 2, 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 };
    }
}

var _this = undefined;
var crop = function (img) {
    var size = Math.min(img.shape[0], img.shape[1]);
    var centerHeight = img.shape[0] / 2;
    var beginHeight = centerHeight - (size / 2);
    var centerWidth = img.shape[1] / 2;
    var beginWidth = centerWidth - (size / 2);
    return img.slice([beginHeight, beginWidth, 0], [size, size, 3]);
};
// convert pixel data into a tensor
var cropAndResizeImage = function (img, dims) {
    if (dims === void 0) { dims = [224, 224]; }
    return __awaiter(_this, void 0, void 0, function () {
        return __generator(this, function (_a) {
            return [2 /*return*/, tidy(function () {
                    var croppedImage = crop(image.resizeBilinear(img, dims));
                    return croppedImage.expandDims(0).toFloat().div(scalar(127)).sub(scalar(1));
                })];
        });
    });
};

var getClasses = function (classes) { return classes.reduce(function (labels, label) {
    var _a;
    if (labels[label] !== undefined) {
        return labels;
    }
    return __assign({}, labels, (_a = {}, _a[label] = Object.keys(labels).length, _a));
}, {}); };

var _this$1 = undefined;
var defaultLayers = function (_a) {
    var classes = _a.classes;
    return [
        layers.flatten({ inputShape: [7, 7, 256] }),
        layers.dense({
            units: 100,
            activation: 'relu',
            kernelInitializer: 'varianceScaling',
            useBias: true
        }),
        layers.dense({
            units: classes,
            kernelInitializer: 'varianceScaling',
            useBias: false,
            activation: 'softmax'
        })
    ];
};
var getBatchSize = function (batchSize, xs) {
    if (batchSize) {
        return batchSize;
    }
    if (xs !== undefined) {
        return Math.floor(xs.shape[0] * 0.4);
    }
    return undefined;
};
var train$1 = function (_a, classes, params) {
    var xs = _a.xs, ys = _a.ys;
    return __awaiter(_this$1, void 0, void 0, function () {
        var model$$1, optimizer, batchSize, history;
        return __generator(this, function (_b) {
            switch (_b.label) {
                case 0:
                    if (xs === undefined || ys === undefined) {
                        throw new Error('Add some examples before training!');
                    }
                    model$$1 = sequential({
                        layers: defaultLayers({ classes: classes }),
                    });
                    optimizer = train.adam(0.0001);
                    model$$1.compile({
                        optimizer: optimizer,
                        loss: 'categoricalCrossentropy',
                        metrics: ['accuracy'],
                    });
                    batchSize = getBatchSize(params.batchSize, xs);
                    return [4 /*yield*/, model$$1.fit(xs, ys, __assign({}, params, { batchSize: batchSize, epochs: params.epochs || 20 }))];
                case 1:
                    history = _b.sent();
                    return [2 /*return*/, {
                            model: model$$1,
                            history: history,
                        }];
            }
        });
    });
};

var _a;
var _this$2 = undefined;
var PRETRAINED_MODELS_KEYS = {
    MOBILENET: 'mobilenet_v1_0.25_224',
};
var PRETRAINED_MODELS = (_a = {},
    _a[PRETRAINED_MODELS_KEYS.MOBILENET] = {
        url: 'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json',
        layer: 'conv_pw_13_relu',
    },
    _a);
var loadPretrainedModel = function (pretrainedModel) { return __awaiter(_this$2, void 0, void 0, function () {
    var config, model$$1, layer;
    return __generator(this, function (_a) {
        switch (_a.label) {
            case 0:
                if (!PRETRAINED_MODELS[pretrainedModel]) {
                    throw new Error('You have supplied an invalid key for a pretrained model');
                }
                config = PRETRAINED_MODELS[pretrainedModel];
                return [4 /*yield*/, loadModel(config.url)];
            case 1:
                model$$1 = _a.sent();
                layer = model$$1.getLayer(config.layer);
                return [2 /*return*/, model({
                        inputs: [model$$1.inputs[0]],
                        outputs: layer.output,
                    })];
        }
    });
}); };

var oneHot$1 = function (labelIndex, classLength) { return tidy(function () { return oneHot(tensor1d([labelIndex]).toInt(), classLength); }); };
// const turnTensorArrayIntoTensor = (tensors: tf.Tensor[]) => tensors.reduce((data?: tf.Tensor, tensor: tf.Tensor) => tf.tidy(() => {
//   if (data === undefined) {
//     return tf.keep(tensor);
//   }
//   const newData = tf.keep(data.concat(tensor, 0));
//   data.dispose();
//   return newData;
// }), undefined);
var addData = function (tensors) {
    var data = keep(tensors[0]);
    return tensors.slice(1).reduce(function (data, tensor) { return tidy(function () {
        var newData = keep(data.concat(tensor, 0));
        data.dispose();
        return newData;
    }); }, data);
};
var addLabels = function (labels, classes) {
    var classLength = Object.keys(classes).length;
    return labels.reduce(function (data, label) {
        var labelIndex = classes[label];
        var y = oneHot$1(labelIndex, classLength);
        return tidy(function () {
            if (data === undefined) {
                return keep(y);
            }
            var old = data;
            var ys = keep(old.concat(y, 0));
            old.dispose();
            y.dispose();
            return ys;
        });
    }, undefined);
};

var getOrderedClasses = function (classes) { return Object.entries(classes).sort(function (a, b) {
    return a[1] - b[1];
}).map(function (_a) {
    var key = _a[0];
    return key;
}); };
var getName = function (classes) { return getOrderedClasses(classes).join('-'); };
var getDefaultDownloadHandler = function (classes) {
    return "downloads://ml-classifier-" + getName(classes);
};

var DataType;
(function (DataType) {
    DataType["TRAIN"] = "train";
    DataType["EVAL"] = "eval";
})(DataType || (DataType = {}));

var MLClassifier = /** @class */ (function () {
    function MLClassifier() {
        var _this = this;
        this.callbacks = [];
        this.data = {
            classes: {},
        };
        this.init = function () { return __awaiter(_this, void 0, void 0, function () {
            var _a;
            return __generator(this, function (_b) {
                switch (_b.label) {
                    case 0:
                        _a = this;
                        return [4 /*yield*/, loadPretrainedModel(PRETRAINED_MODELS_KEYS.MOBILENET)];
                    case 1:
                        _a.pretrainedModel = _b.sent();
                        this.callbacks.map(function (callback) { return callback(); });
                        return [2 /*return*/];
                }
            });
        }); };
        this.loaded = function () { return __awaiter(_this, void 0, void 0, function () {
            var _this = this;
            return __generator(this, function (_a) {
                return [2 /*return*/, new Promise(function (resolve) {
                        if (_this.pretrainedModel) {
                            return resolve();
                        }
                        _this.callbacks.push(function () {
                            resolve();
                        });
                    })];
            });
        }); };
        this.cropAndActivateImage = function (image$$1) { return __awaiter(_this, void 0, void 0, function () {
            var processedImage;
            return __generator(this, function (_a) {
                switch (_a.label) {
                    case 0: return [4 /*yield*/, this.loaded()];
                    case 1:
                        _a.sent();
                        return [4 /*yield*/, cropAndResizeImage(image$$1)];
                    case 2:
                        processedImage = _a.sent();
                        return [2 /*return*/, this.pretrainedModel.predict(processedImage)];
                }
            });
        }); };
        this.getData = function (dataType) { return __awaiter(_this, void 0, void 0, function () {
            return __generator(this, function (_a) {
                if (!this.data[dataType]) {
                    throw new Error("Datatype " + dataType + " unsupported");
                }
                return [2 /*return*/, {
                        xs: this.data[dataType].xs,
                        ys: this.data[dataType].ys,
                        classes: this.data.classes,
                    }];
            });
        }); };
        this.getModel = function () { return _this.model; };
        this.addData = function (images, labels, dataType) {
            if (dataType === void 0) { dataType = DataType.TRAIN; }
            return __awaiter(_this, void 0, void 0, function () {
                var activatedImages, xs, ys;
                var _this = this;
                return __generator(this, function (_a) {
                    switch (_a.label) {
                        case 0:
                            if (!images) {
                                throw new Error('You must supply images');
                            }
                            if (!labels) {
                                throw new Error('You must supply labels');
                            }
                            if (!(dataType === DataType.TRAIN || dataType === DataType.EVAL)) return [3 /*break*/, 2];
                            return [4 /*yield*/, Promise.all(images.map(function (image$$1, idx) { return __awaiter(_this, void 0, void 0, function () {
                                    return __generator(this, function (_a) {
                                        switch (_a.label) {
                                            case 0: return [4 /*yield*/, nextFrame()];
                                            case 1:
                                                _a.sent();
                                                return [4 /*yield*/, this.cropAndActivateImage(image$$1)];
                                            case 2: return [2 /*return*/, _a.sent()];
                                        }
                                    });
                                }); }))];
                        case 1:
                            activatedImages = _a.sent();
                            this.data.classes = getClasses(labels);
                            xs = addData(activatedImages);
                            ys = addLabels(labels, this.data.classes);
                            this.data[dataType] = {
                                xs: xs,
                                ys: ys,
                            };
                            _a.label = 2;
                        case 2: return [2 /*return*/];
                    }
                });
            });
        };
        this.clearData = function (dataType) { return __awaiter(_this, void 0, void 0, function () {
            return __generator(this, function (_a) {
                if (dataType) {
                    this.data[dataType] = {};
                }
                this.data[DataType.TRAIN] = {};
                this.data[DataType.EVAL] = {};
                return [2 /*return*/];
            });
        }); };
        this.train = function (params) {
            if (params === void 0) { params = {}; }
            return __awaiter(_this, void 0, void 0, function () {
                var data, classes, _a, model$$1, history;
                return __generator(this, function (_b) {
                    switch (_b.label) {
                        case 0: return [4 /*yield*/, this.loaded()];
                        case 1:
                            _b.sent();
                            return [4 /*yield*/, this.getData(DataType.TRAIN)];
                        case 2:
                            data = _b.sent();
                            if (!data.xs) {
                                throw new Error('You must add some training examples');
                            }
                            classes = Object.keys(data.classes).length;
                            return [4 /*yield*/, train$1(data, classes, params)];
                        case 3:
                            _a = _b.sent(), model$$1 = _a.model, history = _a.history;
                            this.model = model$$1;
                            return [2 /*return*/, history];
                    }
                });
            });
        };
        this.predict = function (data) { return __awaiter(_this, void 0, void 0, function () {
            var img, predictedClass, classId;
            var _this = this;
            return __generator(this, function (_a) {
                switch (_a.label) {
                    case 0: return [4 /*yield*/, this.loaded()];
                    case 1:
                        _a.sent();
                        if (!this.model) {
                            throw new Error('You must call train prior to calling predict');
                        }
                        return [4 /*yield*/, this.cropAndActivateImage(data)];
                    case 2:
                        img = _a.sent();
                        predictedClass = tidy(function () {
                            var predictions = _this.model.predict(img);
                            // TODO: address this
                            return predictions.as1D().argMax();
                        });
                        return [4 /*yield*/, predictedClass.data()];
                    case 3:
                        classId = (_a.sent())[0];
                        predictedClass.dispose();
                        return [2 /*return*/, Object.entries(this.data.classes).reduce(function (obj, _a) {
                                var key = _a[0], val = _a[1];
                                var _b;
                                return (__assign({}, obj, (_b = {}, _b[val] = key, _b)));
                            }, {})[classId]];
                }
            });
        }); };
        this.evaluate = function (params) {
            if (params === void 0) { params = {}; }
            return __awaiter(_this, void 0, void 0, function () {
                var data;
                return __generator(this, function (_a) {
                    switch (_a.label) {
                        case 0: return [4 /*yield*/, this.loaded()];
                        case 1:
                            _a.sent();
                            if (!this.model) {
                                throw new Error('You must call train prior to calling predict');
                            }
                            return [4 /*yield*/, this.getData(DataType.EVAL)];
                        case 2:
                            data = _a.sent();
                            if (!data.xs || !data.ys) {
                                throw new Error('You must add some evaluation examples');
                            }
                            return [4 /*yield*/, this.model.evaluate(data.xs, data.ys, params)];
                        case 3: return [2 /*return*/, _a.sent()];
                    }
                });
            });
        };
        // handlerOrURL?: tf.io.IOHandler | string;
        this.save = function (handlerOrURL, params) {
            if (params === void 0) { params = {}; }
            return __awaiter(_this, void 0, void 0, function () {
                return __generator(this, function (_a) {
                    switch (_a.label) {
                        case 0: return [4 /*yield*/, this.loaded()];
                        case 1:
                            _a.sent();
                            if (!this.model) {
                                throw new Error('You must call train prior to calling save');
                            }
                            return [4 /*yield*/, this.model.save(handlerOrURL || getDefaultDownloadHandler(this.data.classes), params)];
                        case 2: return [2 /*return*/, _a.sent()];
                    }
                });
            });
        };
        this.init();
    }
    return MLClassifier;
}());

export default MLClassifier;
export { DataType };
