import * as tf from '@tensorflow/tfjs-core';
import { TfjsImageRecognitionBase } from 'tfjs-image-recognition-base';

import { prelu } from './prelu';
import { SharedParams } from './types';

export function sharedLayer(x: tf.Tensor4D, params: SharedParams, isPnet: boolean = false) {
  return tf.tidy(() => {

    let out = TfjsImageRecognitionBase.convLayer(x, params.conv1, 'valid')
    out = prelu<tf.Tensor4D>(out, params.prelu1_alpha)
    out = tf.maxPool(out, isPnet ? [2, 2]: [3, 3], [2, 2], 'same')
    out = TfjsImageRecognitionBase.convLayer(out, params.conv2, 'valid')
    out = prelu<tf.Tensor4D>(out, params.prelu2_alpha)
    out = isPnet ? out : tf.maxPool(out, [3, 3], [2, 2], 'valid')
    out = TfjsImageRecognitionBase.convLayer(out, params.conv3, 'valid')
    out = prelu<tf.Tensor4D>(out, params.prelu3_alpha)

    return out
  })
}