/**
 * BlazeFace, FaceMesh & Iris model implementation
 *
 * Based on:
 * - [**MediaPipe BlazeFace**](https://drive.google.com/file/d/1f39lSzU5Oq-j_OXgS67KfN5wNsoeAZ4V/view)
 * - Facial Spacial Geometry: [**MediaPipe FaceMesh**](https://drive.google.com/file/d/1VFC_wIpw4O7xBOiTgUldl79d9LA-LsnA/view)
 * - Eye Iris Details: [**MediaPipe Iris**](https://drive.google.com/file/d/1bsWbokp9AklH2ANjCfmjqEzzxO1CNbMu/view)
 */

import { log, now } from '../util/util';
import { loadModel } from '../tfjs/load';
import * as tf from '../../dist/tfjs.esm.js';
import * as blazeface from './blazeface';
import * as util from './facemeshutil';
import * as coords from './facemeshcoords';
import * as iris from './iris';
import * as attention from './attention';
import { histogramEqualization } from '../image/enhance';
import { env } from '../util/env';
import type { GraphModel, Tensor } from '../tfjs/types';
import type { FaceResult, FaceLandmark, Point } from '../result';
import type { Config } from '../config';

type DetectBox = { startPoint: Point, endPoint: Point, landmarks: Array<Point>, confidence: number };

const cache = {
  boxes: [] as DetectBox[],
  skipped: Number.MAX_SAFE_INTEGER,
  timestamp: 0,
};

let model: GraphModel | null = null;
let inputSize = 0;

export async function predict(input: Tensor, config: Config): Promise<FaceResult[]> {
  // reset cached boxes
  const skipTime = (config.face.detector?.skipTime || 0) > (now() - cache.timestamp);
  const skipFrame = cache.skipped < (config.face.detector?.skipFrames || 0);
  if (!config.skipAllowed || !skipTime || !skipFrame || cache.boxes.length === 0) {
    cache.boxes = await blazeface.getBoxes(input, config); // get results from blazeface detector
    cache.timestamp = now();
    cache.skipped = 0;
  } else {
    cache.skipped++;
  }
  const faces: Array<FaceResult> = [];
  const newCache: Array<DetectBox> = [];
  let id = 0;
  for (let i = 0; i < cache.boxes.length; i++) {
    const box = cache.boxes[i];
    let angle = 0;
    let rotationMatrix;
    const face: FaceResult = { // init face result
      id: id++,
      mesh: [],
      meshRaw: [],
      box: [0, 0, 0, 0],
      boxRaw: [0, 0, 0, 0],
      score: 0,
      boxScore: 0,
      faceScore: 0,
      // contoursRaw: [],
      // contours: [],
      annotations: {} as Record<FaceLandmark, Point[]>,
    };

    // optional rotation correction based on detector data only if mesh is disabled otherwise perform it later when we have more accurate mesh data. if no rotation correction this function performs crop
    [angle, rotationMatrix, face.tensor] = util.correctFaceRotation(config.face.detector?.rotation, box, input, config.face.mesh?.enabled ? inputSize : blazeface.size());
    if (config?.filter?.equalization) {
      const equilized = await histogramEqualization(face.tensor as Tensor);
      tf.dispose(face.tensor);
      face.tensor = equilized;
    }
    face.boxScore = Math.round(100 * box.confidence) / 100;
    if (!config.face.mesh?.enabled) { // mesh not enabled, return resuts from detector only
      face.box = util.clampBox(box, input);
      face.boxRaw = util.getRawBox(box, input);
      face.score = face.boxScore;
      face.mesh = box.landmarks.map((pt) => [
        ((box.startPoint[0] + box.endPoint[0])) / 2 + ((box.endPoint[0] + box.startPoint[0]) * pt[0] / blazeface.size()),
        ((box.startPoint[1] + box.endPoint[1])) / 2 + ((box.endPoint[1] + box.startPoint[1]) * pt[1] / blazeface.size()),
      ]);
      face.meshRaw = face.mesh.map((pt) => [pt[0] / (input.shape[2] || 0), pt[1] / (input.shape[1] || 0), (pt[2] || 0) / inputSize]);
      for (const key of Object.keys(coords.blazeFaceLandmarks)) {
        face.annotations[key] = [face.mesh[coords.blazeFaceLandmarks[key] as number]]; // add annotations
      }
    } else if (!model) { // mesh enabled, but not loaded
      if (config.debug) log('face mesh detection requested, but model is not loaded');
    } else { // mesh enabled
      if (config.face.attention?.enabled && !env.kernels.includes('atan2')) {
        tf.dispose(face.tensor);
        return faces;
      }
      const results = model.execute(face.tensor as Tensor) as Array<Tensor>;
      const confidenceT = results.find((t) => t.shape[t.shape.length - 1] === 1) as Tensor;
      const faceConfidence = await confidenceT.data();
      face.faceScore = Math.round(100 * faceConfidence[0]) / 100;

      if (face.faceScore < (config.face.detector?.minConfidence || 1)) { // low confidence in detected mesh
        box.confidence = face.faceScore; // reset confidence of cached box
        if (config.face.mesh?.keepInvalid) {
          face.box = util.clampBox(box, input);
          face.boxRaw = util.getRawBox(box, input);
          face.score = face.boxScore;
          face.mesh = box.landmarks.map((pt) => [
            ((box.startPoint[0] + box.endPoint[0])) / 2 + ((box.endPoint[0] + box.startPoint[0]) * pt[0] / blazeface.size()),
            ((box.startPoint[1] + box.endPoint[1])) / 2 + ((box.endPoint[1] + box.startPoint[1]) * pt[1] / blazeface.size()),
          ]);
          face.meshRaw = face.mesh.map((pt) => [pt[0] / (input.shape[2] || 1), pt[1] / (input.shape[1] || 1), (pt[2] || 0) / inputSize]);
          for (const key of Object.keys(coords.blazeFaceLandmarks)) {
            face.annotations[key] = [face.mesh[coords.blazeFaceLandmarks[key] as number]]; // add annotations
          }
        }
      } else {
        const meshT = results.find((t) => t.shape[t.shape.length - 1] === 1404) as Tensor;
        const coordsReshaped = tf.reshape(meshT, [-1, 3]);
        let rawCoords = await coordsReshaped.array();
        tf.dispose(coordsReshaped);
        if (config.face.attention?.enabled) {
          rawCoords = await attention.augment(rawCoords, results); // augment iris results using attention model results
        } else if (config.face.iris?.enabled) {
          rawCoords = await iris.augmentIris(rawCoords, face.tensor, config, inputSize); // run iris model and augment results
        }
        face.mesh = util.transformRawCoords(rawCoords, box, angle, rotationMatrix, inputSize); // get processed mesh
        face.meshRaw = face.mesh.map((pt) => [pt[0] / (input.shape[2] || 0), pt[1] / (input.shape[1] || 0), (pt[2] || 0) / inputSize]);
        for (const key of Object.keys(coords.meshAnnotations)) face.annotations[key] = coords.meshAnnotations[key].map((index) => face.mesh[index]); // add annotations
        face.score = face.faceScore;
        const calculatedBox = { ...util.calculateFaceBox(face.mesh, box), confidence: box.confidence, landmarks: box.landmarks };
        face.box = util.clampBox(calculatedBox, input);
        face.boxRaw = util.getRawBox(calculatedBox, input);
        /*
        const contoursT = results.find((t) => t.shape[t.shape.length - 1] === 266) as Tensor;
        const contoursData = contoursT && await contoursT.data(); // 133 x 2d points
        face.contoursRaw = [];
        for (let j = 0; j < contoursData.length / 2; j++) face.contoursRaw.push([contoursData[2 * j + 0] / inputSize, contoursData[2 * j + 1] / inputSize]);
        face.contours = face.contoursRaw.map((c) => [Math.trunc((input.shape[2] || 1) * c[0]), Math.trunc((input.shape[1] || 1) * c[1])]);
        */
        newCache.push(calculatedBox);
      }
      tf.dispose(results);
    }
    if (face.score > (config.face.detector?.minConfidence || 1)) faces.push(face);
    else tf.dispose(face.tensor);
  }
  cache.boxes = newCache; // reset cache
  return faces;
}

export async function load(config: Config): Promise<GraphModel> {
  if (env.initial) model = null;
  // @ts-ignore private property
  if (config?.face?.attention?.enabled && model?.signature) {
    // @ts-ignore private property
    if (Object.keys(model?.signature?.outputs || {}).length < 6) model = null;
  }
  if (!model) {
    if (config.face.attention?.enabled) model = await loadModel(config.face.attention?.modelPath);
    else model = await loadModel(config.face.mesh?.modelPath);
  } else if (config.debug) {
    log('cached model:', model['modelUrl']);
  }
  inputSize = model.inputs[0].shape ? model.inputs[0].shape[2] : 0;
  return model;
}

export const triangulation = coords.TRI468;
export const uvmap = coords.UV468;