#!/usr/bin/env python3
"""Build NNCF calibration tensors for the inswapper export, from real face photos.

Calibrating a face model with random noise gives useless activation ranges (the
INT8 recogniser embeddings drift by ~0.2 cosine).  This script runs the *reference*
ONNX pipeline on real images, and stores the exact tensors each network sees:

    calib/inswapper_target.npy   (N, 3, 128, 128) float32  RGB in [0, 1]
    calib/inswapper_source.npy   (N, 512)         float32  projected source embedding
    calib/arcface_input.npy      (N, 3, 112, 112) float32  RGB, (x-127.5)/127.5
    calib/retinaface_input.npy   (N, 3, 640, 640) uint8-ish RGB, (x-127.5)/128

Usage:
    python make_calib_data.py --images faces/*.jpg --out-dir calib --samples 24
"""
import argparse
import glob
import os
import sys

import cv2
import numpy as np
import onnxruntime as ort

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(HERE))
import inswapper_ov as iov  # noqa: E402


class ORTSession:
    def __init__(self, path):
        self.sess = ort.InferenceSession(path, providers=['CPUExecutionProvider'])
        self.names = [i.name for i in self.sess.get_inputs()]

    def __call__(self, inputs):
        return self.sess.run(None, {n: inputs[i] for i, n in enumerate(self.names)})


def augment(img, rng, scale_range=(0.75, 1.35), max_rot=8):
    """Random resize + rotation so that a handful of photos give usable statistics."""
    h, w = img.shape[:2]
    scale = rng.uniform(*scale_range)
    ang = rng.uniform(-max_rot, max_rot)
    matrix = cv2.getRotationMatrix2D((w / 2, h / 2), ang, scale)
    cos, sin = abs(matrix[0, 0]), abs(matrix[0, 1])
    nw, nh = int(h * cos + w * sin), int(h * sin + w * cos)
    matrix[0, 2] += nw / 2 - w / 2
    matrix[1, 2] += nh / 2 - h / 2
    return cv2.warpAffine(img, matrix, (nw, nh), borderMode=cv2.BORDER_REFLECT_101)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--images', nargs='+', required=True)
    parser.add_argument('--onnx-dir', default='onnx')
    parser.add_argument('--out-dir', default='calib')
    parser.add_argument('--samples', type=int, default=24)
    args = parser.parse_args()

    paths = []
    for pattern in args.images:
        paths.extend(sorted(glob.glob(pattern)))
    if not paths:
        raise SystemExit('no images given')

    detector = iov.RetinaFace(ORTSession(os.path.join(args.onnx_dir, 'retinaface_10g.onnx')))
    recognizer = iov.ArcFace(ORTSession(os.path.join(args.onnx_dir, 'arcface_w600k_r50.onnx')))
    import onnx
    emap = onnx.numpy_helper.to_array(onnx.load(os.path.join(args.onnx_dir, 'inswapper_128.onnx'))
                                      .graph.initializer[-1]).astype(np.float32)
    swapper = iov.INSwapper(ORTSession(os.path.join(args.onnx_dir, 'inswapper_128.onnx')), emap)

    rng = np.random.default_rng(0)
    swap_target, swap_source, arc_in, det_in = [], [], [], []
    for path in paths:
        base = cv2.imread(path)
        for i in range(max(1, args.samples // len(paths))):
            img = augment(base, rng) if i else base
            _, kpss = detector.detect(img)
            if len(kpss) == 0:
                continue
            kps = kpss[0]
            aimg, _ = iov.norm_crop(img, kps, 128)
            arc = cv2.dnn.blobFromImage(aimg, 1 / 127.5, (112, 112), (127.5, 127.5, 127.5), swapRB=True)
            emb = recognizer.compiled([arc])[0][0].flatten()
            emb = emb / np.linalg.norm(emb)
            latent = np.dot(emb.reshape(1, -1), emap)
            latent = (latent / np.linalg.norm(latent)).astype(np.float32)
            swap_target.append(cv2.dnn.blobFromImage(aimg, 1 / 255.0, (128, 128), (0, 0, 0), swapRB=True))
            swap_source.append(latent)
            arc_in.append(arc)
            h, w = img.shape[:2]
            det_img = np.zeros((640, 640, 3), np.uint8)
            s = min(640 / w, 640 / h)
            r = cv2.resize(img, (int(w * s), int(h * s)))
            det_img[:r.shape[0], :r.shape[1]] = r
            det_in.append(cv2.dnn.blobFromImage(det_img, 1 / 128.0, (640, 640),
                                                (127.5, 127.5, 127.5), swapRB=True))

    os.makedirs(args.out_dir, exist_ok=True)
    for name, data in (('inswapper_target', swap_target), ('inswapper_source', swap_source),
                       ('arcface_input', arc_in), ('retinaface_input', det_in)):
        arr = np.asarray(data, dtype=np.float32)
        path = os.path.join(args.out_dir, f'{name}.npy')
        np.save(path, arr)
        print(f'{path}: {arr.shape} {arr.dtype} min={arr.min():.3f} max={arr.max():.3f}')


if __name__ == '__main__':
    main()