#!/usr/bin/env python3
"""기존 CAM++ extractor와 GPU extractor의 embedding 수치 차이와 latency를 비교한다."""

import argparse
import random
import sys
import time
from pathlib import Path

import librosa
import torch
import torch.nn.functional as F

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from voxcpm.modules.locdit.onnx import CAMPLUS_ONNX_PATH, EmbeddingExtractor, EmbeddingExtractorGPU


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--audio", type=str, default="", help="16 kHz로 로드할 검증용 wav/flac 경로")
    parser.add_argument("--seed", type=int, default=0, help="10초 초과 audio crop 고정 seed")
    parser.add_argument("--warmup", type=int, default=2, help="latency 측정 전 warmup 반복 횟수")
    parser.add_argument("--iters", type=int, default=5, help="latency 측정 반복 횟수")
    parser.add_argument("--max-abs-error", type=float, default=1e-3)
    parser.add_argument("--mean-abs-error", type=float, default=1e-4)
    parser.add_argument("--min-cosine", type=float, default=0.99999)
    return parser.parse_args()


def resolve_audio_path(audio_arg: str) -> Path:
    if audio_arg:
        path = Path(audio_arg)
    else:
        path = PROJECT_ROOT / "seed-tts-eval" / "seedtts_testset" / "en" / "prompt-wavs" / "common_voice_en_2331.wav"
    if not path.exists():
        raise FileNotFoundError(f"Audio file not found: {path}. Pass --audio explicitly.")
    return path


def load_audio(path: Path, seed: int, max_len: int) -> torch.Tensor:
    audio, _ = librosa.load(str(path), sr=16000, mono=True)
    speech = torch.from_numpy(audio).float()
    if speech.numel() > max_len:
        rng = random.Random(seed)
        start = rng.randint(0, speech.numel() - max_len)
        speech = speech[start : start + max_len].contiguous()
    return speech


def timed_call(fn, warmup: int, iters: int):
    for _ in range(warmup):
        fn()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    start = time.perf_counter()
    result = None
    for _ in range(iters):
        result = fn()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    elapsed = (time.perf_counter() - start) / max(1, iters)
    return result, elapsed


def main():
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("GPU verification requires torch.cuda.is_available() == True")
    if not CAMPLUS_ONNX_PATH.exists():
        raise FileNotFoundError(f"CAM++ ONNX model not found: {CAMPLUS_ONNX_PATH}")

    audio_path = resolve_audio_path(args.audio)
    speech = load_audio(audio_path, args.seed, max_len=10 * 16000)
    speech_gpu = speech.to("cuda", non_blocking=True)

    reference_extractor = EmbeddingExtractor(model_path=str(CAMPLUS_ONNX_PATH))
    gpu_extractor = EmbeddingExtractorGPU(model_path=str(CAMPLUS_ONNX_PATH))

    ref_embedding, ref_latency = timed_call(
        lambda: reference_extractor.inference(speech),
        warmup=args.warmup,
        iters=args.iters,
    )
    gpu_embedding, gpu_latency = timed_call(
        lambda: gpu_extractor.inference(speech_gpu),
        warmup=args.warmup,
        iters=args.iters,
    )

    ref = ref_embedding.detach().float().cpu().flatten()
    gpu = gpu_embedding.detach().float().cpu().flatten()
    diff = (ref - gpu).abs()
    max_abs_error = diff.max().item()
    mean_abs_error = diff.mean().item()
    cosine = F.cosine_similarity(ref.unsqueeze(0), gpu.unsqueeze(0)).item()

    print(f"audio: {audio_path}")
    print(f"shape/reference: {tuple(ref.shape)}")
    print(f"shape/gpu:       {tuple(gpu.shape)}")
    print(f"max_abs_error:   {max_abs_error:.8g}")
    print(f"mean_abs_error:  {mean_abs_error:.8g}")
    print(f"cosine:          {cosine:.8g}")
    print(f"latency/ref:     {ref_latency * 1000:.3f} ms")
    print(f"latency/gpu:     {gpu_latency * 1000:.3f} ms")

    passed = (
        ref.shape == gpu.shape == (192,)
        and max_abs_error <= args.max_abs_error
        and mean_abs_error <= args.mean_abs_error
        and cosine >= args.min_cosine
    )
    print("result:          PASS" if passed else "result:          FAIL")
    if not passed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()