"""Easy4K SFDP · Siamese NET AI 推理示例。

用法：
    python example_infer.py [模型文件] [可选：视频文件]

不传视频时用合成数据自检；传视频时对首两帧取同位置块做判重演示。
"""
import sys

import numpy as np
import onnxruntime as ort

SIZE = 128
THRESHOLD = 0.5


def make_session(model_path="model.onnx"):
    """DirectML 可用就用 GPU，否则自动回退 CPU。"""
    avail = ort.get_available_providers()
    providers = [p for p in ("DmlExecutionProvider", "CPUExecutionProvider") if p in avail]
    sess = ort.InferenceSession(model_path, providers=providers)
    print("providers:", sess.get_providers())
    return sess


def prep(patch):
    """灰度块 -> (1, 1, 128, 128) float32，值域 [0, 1]。"""
    x = np.asarray(patch, dtype=np.float32)
    if x.max() > 1.5:          # 传入的是 0~255 时自动归一化
        x = x / 255.0
    if x.ndim == 3:            # 兼容传入灰度单通道 (128,128,1)
        x = x[..., 0]
    return x.reshape(1, 1, SIZE, SIZE)


def score(sess, gray_a, gray_b):
    a, b = prep(gray_a), prep(gray_b)
    return float(sess.run(["score"], {"frame_a": a, "frame_b": b})[0][0, 0])


def score_batch(sess, pairs_a, pairs_b):
    """批量推理：pairs 形状 (N, 128, 128)，返回 (N,) 分数。"""
    a = np.stack([prep(p).reshape(SIZE, SIZE) for p in pairs_a], axis=0).astype(np.float32)
    b = np.stack([prep(p).reshape(SIZE, SIZE) for p in pairs_b], axis=0).astype(np.float32)
    a = a[:, None, :, :]
    b = b[:, None, :, :]
    return sess.run(["score"], {"frame_a": a, "frame_b": b})[0].ravel()


def crop(gray, x, y):
    p = gray[y:y + SIZE, x:x + SIZE]
    return p if p.shape[:2] == (SIZE, SIZE) else None


def main():
    model = sys.argv[1] if len(sys.argv) > 1 else "model.onnx"
    video = sys.argv[2] if len(sys.argv) > 2 else None
    sess = make_session(model)

    if not video:
        base = np.random.randint(0, 255, (SIZE, SIZE), np.uint8)
        print("same  ->", round(score(sess, base, base), 4))
        print("noise ->", round(score(sess, base, np.random.randint(0, 255, (SIZE, SIZE), np.uint8)), 4))
        pairs = [base] * 4
        print("batch ->", np.round(score_batch(sess, pairs, pairs), 4))
        return

    import cv2

    cap = cv2.VideoCapture(video)
    ok1, f1 = cap.read()
    ok2, f2 = cap.read()
    cap.release()
    if not (ok1 and ok2):
        raise SystemExit("视频读取失败")

    g1 = cv2.cvtColor(f1, cv2.COLOR_BGR2GRAY)
    g2 = cv2.cvtColor(f2, cv2.COLOR_BGR2GRAY)

    h, w = g1.shape
    a, b = crop(g1, (w - SIZE) // 2, (h - SIZE) // 2), crop(g2, (w - SIZE) // 2, (h - SIZE) // 2)
    if a is None:
        raise SystemExit("视频分辨率小于 128x128")

    s = score(sess, a, b)
    print("score =", round(s, 4), "| duplicate" if s >= THRESHOLD else "| different")


if __name__ == "__main__":
    main()