import { Mat } from "../../cv/mat";
import { convertImage } from "../common/processors";
import { FaceRecognition, FaceRecognitionPredictOption } from "./common";

const MODEL_URL_CONFIG_ARC_FACE = {
	INSIGHTFACE_ARCFACE_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r100.onnx`,
	INSIGHTFACE_ARCFACE_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r50.onnx`,
	INSIGHTFACE_ARCFACE_R34_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r34.onnx`,
	INSIGHTFACE_ARCFACE_R18_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r18.onnx`,
	INSIGHTFACE_ARCFACE_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r50.mnn`,
	INSIGHTFACE_ARCFACE_R34_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r34.mnn`,
	INSIGHTFACE_ARCFACE_R18_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/ms1mv3_arcface_r18.mnn`,
};
const MODEL_URL_CONFIG_COS_FACE = {
	INSIGHTFACE_COSFACE_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r100.onnx`,
	INSIGHTFACE_COSFACE_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r50.onnx`,
	INSIGHTFACE_COSFACE_R34_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r34.onnx`,
	INSIGHTFACE_COSFACE_R18_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r18.onnx`,
	INSIGHTFACE_COSFACE_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r50.mnn`,
	INSIGHTFACE_COSFACE_R34_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r34.mnn`,
	INSIGHTFACE_COSFACE_R18_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/glint360k_cosface_r18.mnn`,
};
const MODEL_URL_CONFIG_PARTIAL_FC = {
	INSIGHTFACE_PARTIALFC_R100_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r100.onnx`,
	INSIGHTFACE_PARTIALFC_R50_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r50.onnx`,
	INSIGHTFACE_PARTIALFC_R50_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceid/insightface/partial_fc_glint360k_r50.mnn`,
};


export class Insightface extends FaceRecognition {

	public async doPredict(image: Mat, option?: FaceRecognitionPredictOption): Promise<number[]> {
		const input = this.input;
		const output = this.output;
		if (option?.crop) image = image.crop(option.crop.sx, option.crop.sy, option.crop.sw, option.crop.sh);
		image = image.resize(input.shape[3], input.shape[2]);
		const nchwImageData = convertImage(image.data, { sourceImageFormat: "bgr", targetColorFormat: "bgr", targetShapeFormat: "nchw", targetNormalize: { mean: [127.5], std: [127.5] } });

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

		return new Array(...embedding);
	}

}

export class ArcFace extends Insightface {
	public static async load(type?: keyof typeof MODEL_URL_CONFIG_ARC_FACE) {
		return this.cacheModel(MODEL_URL_CONFIG_ARC_FACE[type ?? "INSIGHTFACE_ARCFACE_R100_ONNX"], { createModel: true }).then(r => r.model);
	}
}

export class CosFace extends Insightface {
	public static async load(type?: keyof typeof MODEL_URL_CONFIG_COS_FACE) {
		return this.cacheModel(MODEL_URL_CONFIG_COS_FACE[type ?? "INSIGHTFACE_COSFACE_R100_ONNX"], { createModel: true }).then(r => r.model);
	}
}

export class PartialFC extends Insightface {
	public static async load(type?: keyof typeof MODEL_URL_CONFIG_PARTIAL_FC) {
		return this.cacheModel(MODEL_URL_CONFIG_PARTIAL_FC[type ?? "INSIGHTFACE_PARTIALFC_R100_ONNX"], { createModel: true }).then(r => r.model);
	}
}