import * as tf from '@tensorflow/tfjs';

export function isTensor(c: any): c is tf.Tensor {
  return typeof c.dataId === 'object' && typeof c.shape === 'object';
}

export async function loadModel(): Promise<tf.LayersModel>  {
  const trainLayerV1 = 'conv_pw_13_relu';
  const trainLayerV2 = 'out_relu';

  var mobileNetVersion = 2;
  const modelURL = mobileNetURL(mobileNetVersion);
  const load_model = await tf.loadLayersModel(modelURL);
  
  if (mobileNetVersion == 1) {
    const layer = load_model.getLayer(trainLayerV1);
    const truncatedModel = tf.model({
      inputs: load_model.inputs, 
      outputs: layer.output
    })
    const model = tf.sequential();
    model.add(truncatedModel);
    model.add(tf.layers.flatten());
    return model;
  } else {
    const layer = load_model.getLayer(trainLayerV2);
    const truncatedModel = tf.model({
      inputs: load_model.inputs, 
      outputs: layer.output
    })
    const model = tf.sequential();
    model.add(truncatedModel);
    model.add(tf.layers.globalAveragePooling2d({})); // go from shape [7, 7, 1280] to [1280]
    return model;
  }
}

export function mobileNetURL(version: number): string {
  if (version == 1) {
    return "https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_1.0_224/model.json";
  } 
  return "https://storage.googleapis.com/teachable-machine-models/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224_no_top/model.json";
}

export function imageToTensor(data: any): tf.Tensor3D {
  let tensor: tf.Tensor3D;

  if (data instanceof tf.Tensor) {
      tensor = data;
  } else {
     // MobileNet 모델 로드
      tensor = tf.browser.fromPixels(data);
  }
  return tensor;
}

export async function capture(rasterElement: HTMLImageElement | HTMLVideoElement | HTMLCanvasElement, grayscale?: boolean) {
  return tf.tidy(() => {
    const pixels = tf.browser.fromPixels(rasterElement);

    // crop the image so we're using the center square
    const cropped = cropTensor(pixels, grayscale);

    // Expand the outer most dimension so we have a batch size of 1
    const batchedImage = cropped.expandDims(0);

    // Normalize the image between -1 and a1. The image comes in between 0-255
    // so we divide by 127 and subtract 1.
    return batchedImage.toFloat().div(tf.scalar(127)).sub(tf.scalar(1));
  });
}


export function cropTensor( img: tf.Tensor3D, grayscaleModel?: boolean, grayscaleInput?: boolean ) : tf.Tensor3D {
  const size = Math.min(img.shape[0], img.shape[1]);
  const centerHeight = img.shape[0] / 2;
  const beginHeight = centerHeight - (size / 2);
  const centerWidth = img.shape[1] / 2;
  const beginWidth = centerWidth - (size / 2);
  
  if (grayscaleModel && !grayscaleInput) {
    //cropped rgb data
    let grayscale_cropped = img.slice([beginHeight, beginWidth, 0], [size, size, 3]);
    
    grayscale_cropped = grayscale_cropped.reshape([size * size, 1, 3])
    const rgb_weights = [0.2989, 0.5870, 0.1140]
    grayscale_cropped = tf.mul(grayscale_cropped, rgb_weights)
    grayscale_cropped = grayscale_cropped.reshape([size, size, 3]);

    grayscale_cropped = tf.sum(grayscale_cropped, -1)
    grayscale_cropped = tf.expandDims(grayscale_cropped, -1)
    return grayscale_cropped;
  }
  return img.slice([beginHeight, beginWidth, 0], [size, size, 3]);
}