import { writeFileSync } from "fs";
import { cv } from "../../cv";
import { ImageCropOption, ImageSource, Model } from "../common/model";
import { convertImage } from "../common/processors";
import { FaceAlignmentResult, FacePoint } from "./common";

export interface PFLDPredictOption extends ImageCropOption { }

class PFLDResult extends FaceAlignmentResult {
	protected directionPointIndex(): [number, number] { return [36, 92]; }
	protected leftEyePointIndex(): number[] { return [33, 34, 35, 36, 37, 38, 39, 40, 41, 42]; }
	protected rightEyePointIndex(): number[] { return [87, 88, 89, 90, 91, 92, 93, 94, 95, 96]; }
	protected leftEyebrowPointIndex(): number[] { return [43, 44, 45, 46, 47, 48, 49, 50, 51]; }
	protected rightEyebrowPointIndex(): number[] { return [97, 98, 99, 100, 101, 102, 103, 104, 105]; }
	protected mouthPointIndex(): number[] { return [52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71]; }
	protected nosePointIndex(): number[] { return [72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86]; }
	protected contourPointIndex(): number[] { return [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32]; }
}

export class PFLD extends Model {
	public predict(image: ImageSource, option?: PFLDPredictOption) { return Model.resolveImage(image, im => this.doPredict(im, option)); }

	private async doPredict(image: cv.Mat, option?: PFLDPredictOption) {
		const input = this.input;
		if (option?.crop) image = image.crop(option.crop.sx, option.crop.sy, option.crop.sw, option.crop.sh);
		const ratioWidth = image.width / input.shape[3];
		const ratioHeight = image.height / input.shape[2];
		image = image.resize(input.shape[3], input.shape[2]);

		const nchwImageData = convertImage(image.data, { sourceImageFormat: "bgr", targetColorFormat: "bgr", targetShapeFormat: "nchw", targetNormalize: { mean: [0], std: [255] } })

		const pointsOutput = Object.entries(this.outputs).filter(([_, out]) => out.shape.length == 2)[0][1];

		const res = await this.session.run({
			[input.name]: {
				type: "float32",
				data: nchwImageData,
				shape: [1, 3, input.shape[2], input.shape[3]],
			}
		});
		const pointsBuffer = res[pointsOutput.name];

		const points: FacePoint[] = [];
		for (let i = 0; i < pointsBuffer.length; i += 2) {
			const x = pointsBuffer[i] * image.width * ratioWidth;
			const y = pointsBuffer[i + 1] * image.height * ratioHeight;
			points.push({ x, y });
		}

		return new PFLDResult(points);
	}

}