import fs from "fs";
import { deploy } from "./deploy";
import { cv } from "./cv";
import { faceidTestData } from "./test_data/faceid";
import path from "path";
import crypto from "crypto";
import ai from ".";

async function cacheImage(group: string, url: string) {
	const _url = new URL(url);
	const cacheDir = path.join(__dirname, "../cache/images", group);
	fs.mkdirSync(cacheDir, { recursive: true });
	const cacheJsonFile = path.join(cacheDir, "config.json");
	let jsonData: Record<string, string> = {};
	if (cacheJsonFile && fs.existsSync(cacheJsonFile)) {
		jsonData = JSON.parse(fs.readFileSync(cacheJsonFile, "utf-8"));
		const filename = jsonData[url];
		if (filename && fs.existsSync(filename)) return path.join(cacheDir, filename);
	}
	const data = await fetch(_url).then(res => res.arrayBuffer()).then(buf => new Uint8Array(buf));
	const allowedExtnames = [".jpg", ".jpeg", ".png", ".webp"];
	let extname = path.extname(_url.pathname);
	if (!allowedExtnames.includes(extname)) extname = ".jpeg";
	const md5 = crypto.hash("md5", data, "hex");
	const savename = md5 + extname;
	jsonData[url] = savename;
	fs.writeFileSync(cacheJsonFile, JSON.stringify(jsonData));
	const savepath = path.join(cacheDir, savename);
	fs.writeFileSync(savepath, data);
	return savepath;
}

async function testGenderTest() {
	const facedet = await deploy.facedet.Yolov5Face.load("YOLOV5S_MNN");
	const detector = await deploy.faceattr.GenderAgeDetector.load("INSIGHT_GENDER_AGE_MNN");

	const image = await cv.Mat.load("https://b0.bdstatic.com/ugc/iHBWUj0XqytakT1ogBfBJwc7c305331d2cf904b9fb3d8dd3ed84f5.jpg");
	const boxes = await facedet.predict(image);
	if (!boxes.length) return console.error("未检测到人脸");
	for (const [idx, box] of boxes.entries()) {
		const res = await detector.predict(image, { crop: { sx: box.left, sy: box.top, sw: box.width, sh: box.height } });
		console.log(`[${idx + 1}]`, res);
	}

}

async function testFaceID() {
	console.log("初始化模型")
	const facedet = await deploy.facedet.Yolov5Face.load("YOLOV5S_MNN");
	const faceid = await deploy.faceid.CosFace.load("INSIGHTFACE_COSFACE_R50_MNN");
	const facealign = await deploy.facealign.PFLD.load("PFLD_106_LITE_MNN");
	console.log("初始化模型完成")

	const { basic, tests } = faceidTestData.stars;

	console.log("正在加载图片资源");
	const basicImage = await cv.Mat.load(await cacheImage("faceid", basic.image));
	const testsImages: Record<string, cv.Mat[]> = {};
	for (const [name, imgs] of Object.entries(tests)) {
		testsImages[name] = await Promise.all(imgs.map(img => cacheImage("faceid", img).then(img => cv.Mat.load(img))));
	}

	console.log("正在检测基本数据");
	const basicDetectedFaces = await facedet.predict(basicImage);
	const basicFaceIndex: Record<string, number> = {};
	for (const [name, [x, y]] of Object.entries(basic.faces)) {
		basicFaceIndex[name] = basicDetectedFaces.findIndex(box => box.x1 < x && box.x2 > x && box.y1 < y && box.y2 > y);
	}

	async function getEmbd(image: cv.Mat, box: deploy.facedet.FaceBox) {
		box = box.toSquare();
		const alignResult = await facealign.predict(image, { crop: { sx: box.left, sy: box.top, sw: box.width, sh: box.height } });
		let faceImage = image.rotate(box.centerX, box.centerY, -alignResult.direction * 180 / Math.PI);
		return faceid.predict(faceImage, { crop: { sx: box.left, sy: box.top, sw: box.width, sh: box.height } });
	}

	const basicEmbds: number[][] = [];
	for (const box of basicDetectedFaces) {
		// const embd = await faceid.predict(basicImage, { crop: { sx: box.left, sy: box.top, sw: box.width, sh: box.height } });
		const embd = await getEmbd(basicImage, box);
		basicEmbds.push(embd);
	}

	console.log("正在进行人脸对比");
	for (const [name, image] of Object.entries(testsImages)) {
		console.log(`[${name}] 正在检测`);
		const index = basicFaceIndex[name];
		if (index < 0) {
			console.error(`[${name}] 不存在`);
			continue;
		}
		const basicEmbd = basicEmbds[index]

		for (const [idx, img] of image.entries()) {
			const box = await facedet.predict(img).then(boxes => boxes[0]);
			if (!box) {
				console.error(`[${idx + 1}] 未检测到人脸`);
				continue
			}

			// const embd = await faceid.predict(img, { crop: { sx: box.left, sy: box.top, sw: box.width, sh: box.height } });
			const embd = await getEmbd(img, box);

			const compareEmbds = basicEmbds.map(e => deploy.faceid.cosineDistance(e, embd));
			const max = Math.max(...compareEmbds);
			const min = Math.min(...compareEmbds);

			const distance = deploy.faceid.cosineDistance(basicEmbd, embd);
			console.log(`[${idx + 1}] [${(distance.toFixed(4) == max.toFixed(4)) ? '\x1b[102m成功' : '\x1b[101m失败'}\x1b[0m] 相似度:${distance.toFixed(4)}, 最大:${max.toFixed(4)}, 最小:${min.toFixed(4)}`);
		}
	}

}

async function testFaceAlign() {
	const fd = await deploy.facedet.Yolov5Face.load("YOLOV5S_MNN");
	const fa = await deploy.facealign.PFLD.load("PFLD_106_LITE_MNN");
	// const fa = await deploy.facealign.FaceLandmark1000.load("FACELANDMARK1000_ONNX");
	let image = await cv.Mat.load("https://i0.hdslb.com/bfs/archive/64e47ec9fdac9e24bc2b49b5aaad5560da1bfe3e.jpg");
	image = image.rotate(image.width / 2, image.height / 2, 0);

	const face = await fd.predict(image).then(res => {
		console.log(res);
		return res[0].toSquare()
	});
	const points = await fa.predict(image, { crop: { sx: face.left, sy: face.top, sw: face.width, sh: face.height } });

	points.points.forEach((point, idx) => {
		image.circle(face.left + point.x, face.top + point.y, 2);
	})
	// const point = points.getPointsOf("rightEye")[1];
	// image.circle(face.left + point.x, face.top + point.y, 2);
	fs.writeFileSync("testdata/xx.jpg", image.encode(".jpg"));

	let faceImage = image.rotate(face.centerX, face.centerY, -points.direction * 180 / Math.PI);
	faceImage = faceImage.crop(face.left, face.top, face.width, face.height);
	fs.writeFileSync("testdata/face.jpg", faceImage.encode(".jpg"));

	console.log(points);
	console.log(points.direction * 180 / Math.PI);
}

async function test() {
	// await testGenderTest();
	// await testFaceID();
	await testFaceAlign();
}

test().catch(err => {
	console.error(err);
});