#!/usr/bin/env python3
"""
噪声注入批量测试脚本

功能：
1. 对干净音频按指定 SNR 注入噪声，生成测试样本
2. 调用 C++ 离线推理程序批量检测
3. 输出 CSV 统计结果，便于绘制 Pd-SNR / 定位误差-SNR 曲线

用法示例：
    # 单通道分类测试
    python inject_noise_eval.py \\
        --clean_dir dataset/real/threat \\
        --noise noise_samples/drone_hover.wav \\
        --snrs 20 10 5 0 -5 \\
        --exe build/demo_offline.exe \\
        --output results/threat_vs_drone_noise.csv

    # 多通道定位测试
    python inject_noise_eval.py \\
        --clean_dir dataset/multichannel \\
        --noise noise_samples/white.wav \\
        --snrs 20 10 5 0 \\
        --exe build/demo_offline_multichannel.exe \\
        --num_mics 4 --layout cross \\
        --output results/loc_snr_sweep.csv
"""

import argparse
import csv
import os
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf


def list_wav_files(directory: str):
    """递归列出目录下所有 wav 文件。"""
    return sorted([str(p) for p in Path(directory).rglob("*.wav")])


def compute_rms(signal: np.ndarray):
    """计算信号 RMS。"""
    return np.sqrt(np.mean(signal ** 2))


def mix_with_noise(clean: np.ndarray, noise: np.ndarray, target_snr_db: float):
    """
    将噪声以指定 SNR 混入干净信号。
    如果噪声比干净信号短，会循环拼接。
    支持多通道 clean（noise 会先混为单通道再扩展到匹配通道数，或直接用同通道数）。
    """
    # Ensure 2D: [channels, samples]
    if clean.ndim == 1:
        clean = clean.reshape(1, -1)
    if noise.ndim == 1:
        noise = noise.reshape(1, -1)

    n_ch, n_samples = clean.shape

    # Tile or truncate noise to match length
    noise_tiled = np.tile(noise, (1, int(np.ceil(n_samples / noise.shape[1])) + 1))
    noise_tiled = noise_tiled[:, :n_samples]

    # If noise has fewer channels, replicate first channel
    if noise_tiled.shape[0] < n_ch:
        noise_tiled = np.repeat(noise_tiled[:1, :], n_ch, axis=0)
    elif noise_tiled.shape[0] > n_ch:
        noise_tiled = noise_tiled[:n_ch, :]

    # Compute per-channel RMS and average for a global scaling factor
    clean_rms = compute_rms(clean)
    if clean_rms < 1e-10:
        return clean  # silent signal, return as-is

    noise_rms = compute_rms(noise_tiled)
    if noise_rms < 1e-10:
        return clean

    target_noise_rms = clean_rms / (10 ** (target_snr_db / 20.0))
    scale = target_noise_rms / noise_rms
    mixed = clean + noise_tiled * scale

    # Normalize to prevent clipping
    max_val = np.max(np.abs(mixed))
    if max_val > 1.0:
        mixed = mixed / max_val * 0.95

    return mixed


def run_classifier(exe_path: str, wav_path: str, extra_args: list) -> dict:
    """
    调用 C++ 离线推理程序，解析其标准输出。
    返回字典包含 pred_label, confidence 等。
    """
    import re
    cmd = [exe_path, wav_path] + extra_args
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=30)
        # Decode with GBK first (Windows console default), fallback to UTF-8
        try:
            output = result.stdout.decode("gbk", errors="replace")
        except Exception:
            output = result.stdout.decode("utf-8", errors="replace")
    except Exception as e:
        return {"error": str(e)}

    parsed = {}
    # Robust regex parsing for output like:
    # File: xxx.wav | True: threat | Pred: threat | Conf: 0.9234 | Az: 90.00° | El: 0.00° | Dist: 95.20m
    for line in output.splitlines():
        if line.startswith("File:") and "Pred:" in line:
            m = re.search(r"Pred:\s*(\S+)", line)
            if m:
                parsed["pred"] = m.group(1)
            m = re.search(r"Conf:\s*([0-9.]+)", line)
            if m:
                parsed["conf"] = float(m.group(1))
            m = re.search(r"Az:\s*([0-9.+-]+)", line)
            if m:
                parsed["azimuth"] = float(m.group(1))
            m = re.search(r"El:\s*([0-9.+-]+)", line)
            if m:
                parsed["elevation"] = float(m.group(1))
            m = re.search(r"Dist:\s*([0-9.+-]+)", line)
            if m:
                parsed["distance"] = float(m.group(1))
            break
    return parsed


def main():
    parser = argparse.ArgumentParser(description="Noise injection evaluation for acoustic analyzer")
    parser.add_argument("--clean_dir", required=True, help="Directory containing clean WAV files")
    parser.add_argument("--noise", required=True, help="Noise WAV file path")
    parser.add_argument("--snrs", nargs="+", type=float, required=True,
                        help="List of SNR values in dB, e.g. 20 10 5 0 -5")
    parser.add_argument("--exe", required=True, help="Path to demo_offline or demo_offline_multichannel exe")
    parser.add_argument("--output", required=True, help="Output CSV path")
    parser.add_argument("--num_mics", type=int, default=1, help="Number of channels (for multichannel exe)")
    parser.add_argument("--layout", default="cross", help="Array layout (for multichannel exe)")
    parser.add_argument("--threshold", type=float, default=0.5, help="Detection threshold")
    parser.add_argument("--ground_azimuth", type=float, default=None, help="Ground-truth azimuth for error calc")
    parser.add_argument("--ground_distance", type=float, default=None, help="Ground-truth distance for error calc")
    args = parser.parse_args()

    clean_files = list_wav_files(args.clean_dir)
    if not clean_files:
        print(f"[ERROR] No WAV files found in {args.clean_dir}")
        sys.exit(1)

    noise_data, noise_sr = sf.read(args.noise, dtype="float32")
    if noise_data.ndim == 1:
        noise_data = noise_data.reshape(1, -1)
    else:
        noise_data = noise_data.T  # [channels, samples]

    # Prepare extra args for C++ exe
    extra_args = [f"--threshold", str(args.threshold)]
    if "multichannel" in os.path.basename(args.exe).lower():
        extra_args += ["--num_mics", str(args.num_mics), "--layout", args.layout]
        if args.ground_azimuth is not None:
            extra_args += ["--ground_azimuth", str(args.ground_azimuth)]
        if args.ground_distance is not None:
            extra_args += ["--ground_distance", str(args.ground_distance)]

    # CSV header
    fieldnames = ["file", "true_label", "snr_db", "pred_label", "confidence",
                  "azimuth", "elevation", "distance"]
    if args.ground_azimuth is not None:
        fieldnames.append("azimuth_error")
    if args.ground_distance is not None:
        fieldnames.append("distance_error")

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    with open(args.output, "w", newline="", encoding="utf-8") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        total_runs = len(clean_files) * len(args.snrs)
        run_idx = 0

        with tempfile.TemporaryDirectory() as tmpdir:
            for clean_path in clean_files:
                clean_data, clean_sr = sf.read(clean_path, dtype="float32")
                true_label = Path(clean_path).parent.name

                # Ensure 2D
                if clean_data.ndim == 1:
                    clean_data = clean_data.reshape(1, -1)
                else:
                    clean_data = clean_data.T

                for snr in args.snrs:
                    run_idx += 1
                    mixed = mix_with_noise(clean_data, noise_data, snr)

                    # Save temp WAV (interleaved)
                    if mixed.shape[0] == 1:
                        mixed_interleaved = mixed[0]
                    else:
                        mixed_interleaved = mixed.T  # [samples, channels]

                    tmp_wav = os.path.join(tmpdir, f"mix_{run_idx:04d}.wav")
                    sf.write(tmp_wav, mixed_interleaved, clean_sr, subtype="PCM_16")

                    # Run inference
                    parsed = run_classifier(args.exe, tmp_wav, extra_args)

                    row = {
                        "file": Path(clean_path).name,
                        "true_label": true_label,
                        "snr_db": snr,
                        "pred_label": parsed.get("pred", "error"),
                        "confidence": parsed.get("conf", 0.0),
                        "azimuth": parsed.get("azimuth", ""),
                        "elevation": parsed.get("elevation", ""),
                        "distance": parsed.get("distance", ""),
                    }

                    if args.ground_azimuth is not None and "azimuth" in parsed:
                        err = abs(parsed["azimuth"] - args.ground_azimuth)
                        if err > 180:
                            err = 360 - err
                        row["azimuth_error"] = round(err, 2)

                    if args.ground_distance is not None and "distance" in parsed:
                        row["distance_error"] = round(abs(parsed["distance"] - args.ground_distance), 2)

                    writer.writerow(row)
                    print(f"[{run_idx}/{total_runs}] {Path(clean_path).name} SNR={snr}dB -> {row['pred_label']} ({row['confidence']:.3f})")

    print(f"\n[OK] Results saved to {args.output}")


if __name__ == "__main__":
    main()
