import { encode, decode } from 'base64-arraybuffer'

import { getBackend } from './tf-singleton'

const EstimatorList = [

'KNeighborsRegressor',

'LinearRegression',

'LassoRegression',

'RidgeRegression',

'ElasticNet',

'LogisticRegression',

'DummyRegressor',

'DummyClassifier',

'MinMaxScaler',

'StandardScaler',

'MaxAbsScaler',

'SimpleImputer',

'OneHotEncoder',

'LabelEncoder',

'OrdinalEncoder',

'Normalizer',

'Pipeline',

'ColumnTransformer',

'RobustScaler',

'KMeans',

'VotingRegressor',

'VotingClassifier',

'LinearSVC',

'LinearSVR',

'GaussianNB',

'DecisionTreeClassifier',

'DecisionTreeRegressor',

'ClassificationCriterion',

'RegressionCriterion',

'Splitter',

'DecisionTreeBase',

'DecisionTree'

]

let letters = 'abcdefghijklmnopqrstuvwxy'

function randomString(numLetters: number) {

let curLetter = ''

for (let i = 0; i < numLetters; i++) {

let index = Math.floor(Math.random() * letters.length)

curLetter += letters[index]

}

return curLetter

}

/**

* 1. Make a list called EstimatorList

* 2. Do a dynamic import here

*/

class JSONHandler {

savedArtifacts: any

constructor(artifacts?: any) {

this.savedArtifacts = artifacts || null

}

async save(artifacts: any) {

// Base 64 encoding

artifacts.weightData = encode(artifacts.weightData)

// Remaps the names of the layers, so that when we deserialize we

// don't run into a tfjs error where it says "you've already created these

// names in our backend"

let mapping: any = {}

for (let i = 0; i < artifacts.modelTopology.config.layers.length; i++) {

let curWeightSpec = artifacts.modelTopology.config.layers[i]

let randomName = randomString(6)

mapping[curWeightSpec.config.name] = randomName

curWeightSpec.config.name = randomName

}

for (let i = 0; i < artifacts.weightSpecs.length; i++) {

let cur = artifacts.weightSpecs[i]

let allMaps = Object.keys(mapping)

allMaps.forEach((el) => {

if (cur.name.includes(el)) {

cur.name = cur.name.replace(el, mapping[el])

}

})

}

this.savedArtifacts = artifacts

return {

modelArtifactsInfo: {

dateSaved: new Date(),

modelTopologyType: 'JSON',

modelTopologyBytes: JSON.stringify(artifacts.modelTopology).length,

weightSpecsBytes: JSON.stringify(artifacts.weightSpecs).length,

weightDataBytes: artifacts.weightData.byteLength

}

}

}

async load() {

// Base64 decode

this.savedArtifacts.weightData = decode(this.savedArtifacts.weightData)

return this.savedArtifacts

}

}

export async function toObjectInner(

val: any,

ignoreKeys: string[] = []

): Promise<any> {

let tf = getBackend()

if (['number', 'string', 'undefined', 'boolean'].includes(typeof val)) {

return val

}

if (typeof val === 'function') {

console.warn(

`warning: Serializing function ${val}. Not going to be able to deserialize this later.`

)

if (val.name) {

return val.name

}

}

if (typeof val === 'object') {

// Null case

if (val === null) {

return null

}

// Array case

if (Array.isArray(val)) {

return await Promise.all(

val.map(async (el) => await toObjectInner(el, ignoreKeys))

)

}

// Serialize a Tensor

if (val instanceof tf.Tensor) {

return {

name: 'Tensor',

value: val.arraySync()

}

}

// Int32Array serialization. Used for DecisionTrees

if (val instanceof Int32Array) {

return {

name: 'Int32Array',

value: Array.from(val)

}

}

if (val instanceof Float32Array) {

return {

name: 'Float32Array',

value: Array.from(val)

}

}

// tf.layers model

if (val instanceof tf.Sequential) {

let mem = new JSONHandler()

await val.save(mem as any)

return {

name: 'Sequential',

artifacts: mem.savedArtifacts

}

}

// The tf object

if (val.ENV && val.AdadeltaOptimizer && val.version) {

return {

name: 'TF',

version: val.version.tfjs

}

}

// Generic object case / class case

let response: any = {}

for (let key of Object.keys(val)) {

// Ignore all the keys that we choose to

if (ignoreKeys.includes(key)) {

continue

}

response[key] = await toObjectInner(val[key], ignoreKeys)

}

return response

}

}

export async function fromObjectInner(val: any): Promise<any> {

let tf = getBackend()

// Ignores all types that aren't objects

if (typeof val !== 'object') {

return val

}

// Null case

if (val === null) {

return null

}

// Make a Tensor

if (val.name === 'Tensor') {

return tf.tensor(val.value)

}

if (val.name === 'Sequential') {

let newMem = new JSONHandler(val.artifacts)

return await tf.loadLayersModel(newMem as any)

}

if (val.name === 'Int32Array') {

return new Int32Array(val.value)

}

if (val.name === 'Float32Array') {

return new Float32Array(val.value)

}

if (val.name === 'TF') {

return tf

}

// Array case

if (Array.isArray(val)) {

return await Promise.all(val.map(async (el) => await fromObjectInner(el)))

}

// Generic object case

for (let key of Object.keys(val)) {

val[key] = await fromObjectInner(val[key])

}

// Make a model

if (EstimatorList.includes(val.name)) {

// Do dynamic import to avoid circular dependency tree

// Every class extends this class and therefor it

// can't import those classes in here

let module = await import('./index')

let model = (module as any)[val.name]

let resultObj = new model(val)

for (let key of Object.keys(val)) {

resultObj[key] = val[key]

}

return resultObj

}

return val

}

export async function fromObject(val: any): Promise<any> {

try {

return await fromObjectInner(val)

} catch (e) {

console.error(e)

}

}

export async function fromJSON(val: string): Promise<any> {

return await fromObject(JSON.parse(val))

}

let ignoredKeysForSGDRegressor = [

'modelCompileArgs',

'modelFitArgs',

'denseLayerArgs'

]

export class Serialize {

async toObject(): Promise<any> {

// try {

return await toObjectInner(this, ignoredKeysForSGDRegressor)

// } catch (e) {

// console.error(e)

// }

}

async toJSON(): Promise<string> {

return JSON.stringify(await this.toObject())

}

}