#!/usr/bin/env python3
"""Deterministic CosyVoice peak-memory and reference-quality benchmark."""

import argparse
import math
import os
import sys
import tempfile
import threading
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

PROMPT = "你好，我是会说话的花朵。"
SEED = 20260611
SAMPLE_RATE = 24000
N_FFT = 1024
HOP_LENGTH = 256
N_MELS = 80

_peak_rss_kb = 0
_stop_monitor = False


def _read_rss_kb() -> int:
    try:
        with open("/proc/self/status", "r", encoding="utf-8") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1])
    except OSError:
        return 0
    return 0


def _monitor_rss() -> None:
    global _peak_rss_kb
    while not _stop_monitor:
        rss = _read_rss_kb()
        if rss > _peak_rss_kb:
            _peak_rss_kb = rss
        time.sleep(0.02)


def _start_monitor() -> threading.Thread:
    thread = threading.Thread(target=_monitor_rss, daemon=True)
    thread.start()
    return thread


def _set_deterministic(torch):
    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True


def _synthesize(output_path: Path):
    import torch
    import soundfile as sf
    from cosyvoice.utils.common import set_all_random_seed
    import inference

    set_all_random_seed(SEED)
    _set_deterministic(torch)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

    inference.synthesize(PROMPT, str(output_path))

    if torch.cuda.is_available():
        torch.cuda.synchronize()

    audio, sr = sf.read(str(output_path), dtype="float32")
    if audio.ndim > 1:
        audio = audio[:, 0]
    return audio, sr


def _load_audio(path: Path):
    import soundfile as sf

    audio, sr = sf.read(str(path), dtype="float32")
    if audio.ndim > 1:
        audio = audio[:, 0]
    return audio, sr


def _rms_normalize(audio, eps=1e-8):
    import numpy as np

    rms = float(np.sqrt(np.mean(np.square(audio), dtype=np.float64)))
    if rms < eps:
        return audio
    return audio / rms


def _log_mel(audio, sr):
    import torch
    import torchaudio

    waveform = torch.from_numpy(audio).float().unsqueeze(0)
    transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=sr,
        n_fft=N_FFT,
        hop_length=HOP_LENGTH,
        n_mels=N_MELS,
        center=True,
        power=2.0,
    )
    mel = transform(waveform).clamp_min(1e-8).log()
    return mel.squeeze(0)


def _quality_metrics(candidate, reference, sr):
    import numpy as np
    import torch

    length = min(len(candidate), len(reference))
    candidate = _rms_normalize(candidate[:length])
    reference = _rms_normalize(reference[:length])

    cand_mel = _log_mel(candidate, sr)
    ref_mel = _log_mel(reference, sr)
    frames = min(cand_mel.shape[1], ref_mel.shape[1])
    log_mel_mse = torch.mean(torch.square(cand_mel[:, :frames] - ref_mel[:, :frames])).item()

    err = candidate - reference
    signal_power = float(np.sum(np.square(reference), dtype=np.float64))
    noise_power = float(np.sum(np.square(err), dtype=np.float64)) + 1e-12
    snr_db = 10.0 * math.log10((signal_power + 1e-12) / noise_power)

    duration_s = length / sr
    length_delta_ms = abs(len(candidate) - len(reference)) / sr * 1000.0
    return {
        "log_mel_mse": log_mel_mse,
        "snr_db": snr_db,
        "duration_s": duration_s,
        "length_delta_ms": length_delta_ms,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--reference", default=str(ROOT / "benchmarks" / "reference_fp16.wav"))
    parser.add_argument("--output", default=None)
    parser.add_argument("--write-reference", action="store_true")
    args = parser.parse_args()

    reference_path = Path(args.reference)
    output_path = Path(args.output) if args.output else Path(tempfile.gettempdir()) / "cosyvoice_autoresearch.wav"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    monitor = _start_monitor()
    started = time.perf_counter()
    try:
        audio, sr = _synthesize(output_path)
        if sr != SAMPLE_RATE:
            raise RuntimeError(f"Expected {SAMPLE_RATE} Hz output, got {sr} Hz")

        if args.write_reference:
            reference_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.replace(reference_path)
            print(f"Wrote reference {reference_path}")
            return 0

        if not reference_path.exists():
            raise FileNotFoundError(f"Missing reference audio: {reference_path}")

        reference, ref_sr = _load_audio(reference_path)
        if ref_sr != sr:
            raise RuntimeError(f"Reference sample rate {ref_sr} does not match candidate {sr}")

        quality = _quality_metrics(audio, reference, sr)
        elapsed_s = time.perf_counter() - started

        try:
            import torch
            cuda_peak_mb = torch.cuda.max_memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0
        except Exception:
            cuda_peak_mb = 0.0

        print(f"METRIC peak_rss_mb={_peak_rss_kb / 1024:.3f}")
        print(f"METRIC cuda_peak_allocated_mb={cuda_peak_mb:.3f}")
        print(f"METRIC log_mel_mse={quality['log_mel_mse']:.9f}")
        print(f"METRIC snr_db={quality['snr_db']:.3f}")
        print(f"METRIC length_delta_ms={quality['length_delta_ms']:.3f}")
        print(f"METRIC duration_s={quality['duration_s']:.3f}")
        print(f"METRIC elapsed_s={elapsed_s:.3f}")
        return 0
    finally:
        global _stop_monitor
        _stop_monitor = True
        monitor.join(timeout=0.2)


if __name__ == "__main__":
    raise SystemExit(main())