#!/usr/bin/env python3
"""
验证集准确率验证
遍历 val 目录，使用 ONNX Runtime 进行推理，统计准确率
"""
import argparse
import os
import json
import numpy as np
import onnxruntime as ort
import torchaudio
import torchaudio.functional as F
import librosa


def extract_mel(wav_path, sr=16000, n_fft=2048, hop_length=512, n_mels=64, fmax=8000):
    y, _ = librosa.load(wav_path, sr=sr, mono=True)
    mel = librosa.feature.melspectrogram(
        y=y, sr=sr, n_fft=n_fft, hop_length=hop_length,
        n_mels=n_mels, fmax=fmax, htk=True, center=False
    )
    mel_db = librosa.power_to_db(mel, ref=1.0)
    mel_db = (mel_db - mel_db.mean(axis=1, keepdims=True)) / (mel_db.std(axis=1, keepdims=True) + 1e-8)
    if mel_db.shape[1] < 63:
        mel_db = np.pad(mel_db, ((0, 0), (0, 63 - mel_db.shape[1])), mode='edge')
    else:
        mel_db = mel_db[:, :63]
    return mel_db.astype(np.float32).T  # [T, n_mels]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='models/gunshot_classifier.onnx')
    parser.add_argument('--val_dir', default='dataset/val')
    parser.add_argument('--label_map', default='models/label_map.json')
    parser.add_argument('--threshold', type=float, default=0.7)
    args = parser.parse_args()

    with open(args.label_map) as f:
        label_map = json.load(f)
    idx_to_label = {int(k): v for k, v in label_map.items()}

    session = ort.InferenceSession(args.model, providers=['CPUExecutionProvider'])
    input_name = session.get_inputs()[0].name

    correct = 0
    total = 0
    for true_label in sorted(os.listdir(args.val_dir)):
        cls_dir = os.path.join(args.val_dir, true_label)
        if not os.path.isdir(cls_dir):
            continue
        for fname in sorted(os.listdir(cls_dir)):
            if not fname.endswith('.wav'):
                continue
            path = os.path.join(cls_dir, fname)
            mel = extract_mel(path)
            mel = np.expand_dims(mel, axis=0)  # [1, T, n_mels]
            outputs = session.run(None, {input_name: mel})
            probs = outputs[0][0]
            pred_idx = int(np.argmax(probs))
            pred_conf = float(probs[pred_idx])
            pred_label = idx_to_label.get(pred_idx, 'unknown')
            if pred_conf < args.threshold:
                pred_label = 'ambient'
            if pred_label == true_label:
                correct += 1
            total += 1
            print(f"{path}: pred={pred_label} (conf={pred_conf:.4f}), true={true_label}")

    acc = correct / total if total > 0 else 0.0
    print(f"\nTotal: {total}, Correct: {correct}, Accuracy: {acc:.4f} ({acc*100:.2f}%)")


if __name__ == '__main__':
    main()
