import * as tf from '@tensorflow/tfjs'

type Shape = tf.Shape

export type LayerType =
  | 'conv2d'
  | 'maxPooling2d'
  | 'flatten'
  | 'dense'
  | 'averagePooling2d'
type ActivationType = 'tanh' | 'sigmoid' | 'softmax' | 'relu'

export const activations: ActivationType[] = ['tanh', 'sigmoid', 'softmax', 'relu']

export abstract class Layer {
  layerName: string
  layerType: LayerType

  static conv2dLayer() {
    return new Conv2dLayer()
  }

  static maxPooling2dLayer() {
    return new PoolingLayer('max pooling 2d layer', 'maxPooling2d')
  }

  static averagePooling2dLayer() {
    return new PoolingLayer('average pooling 2d layer', 'averagePooling2d')
  }

  static flattenLayer() {
    return new FalttenLayer()
  }

  static denseLayer() {
    return new DensLayer()
  }

  protected constructor(
    layerName: string,
    layerType: LayerType
  ) {
    this.layerName = layerName
    this.layerType = layerType
  }

  abstract toTfLayer(inputShape?: Shape): tf.layers.Layer
}

class Conv2dLayer extends Layer {
  filters: number = 32
  kernelSize: number = 3
  activation: ActivationType = 'sigmoid'

  constructor() {
    super('conv2d layer', 'conv2d')
  }

  toTfLayer(inputShape?: Shape) {
    return tf.layers.conv2d({
      filters: this.filters,
      kernelSize: this.kernelSize,
      activation: this.activation,
      inputShape
    })
  }
}

class FalttenLayer extends Layer {
  constructor() {
    super('flatten layer', 'flatten')
  }

  toTfLayer(inputShape?: Shape) {
    return tf.layers.flatten({ inputShape })
  }
}

class PoolingLayer extends Layer {
  poolsize: number = 2
  strides: number = 2

  toTfLayer(inputShape?: Shape) {
    if (this.layerType === 'maxPooling2d') {
      return tf.layers.maxPool2d({
        poolSize: this.poolsize,
        strides: this.strides,
        inputShape
      })
    } else {
      return tf.layers.averagePooling2d({
        poolSize: this.poolsize,
        strides: this.strides,
        inputShape
      })
    }
  }
}

class DensLayer extends Layer {
  units: number = 10
  activation: ActivationType = 'softmax'
  constructor() {
    super('dense layer', 'dense')
  }

  toTfLayer(inputShape?: Shape) {
    this.units = Number.parseInt('' + this.units)
    return tf.layers.dense({
      units: this.units,
      activation: this.activation,
      inputShape
    })
  }
}