#!/usr/bin/env python3
"""
C++/Python 特征一致性验证
对比 C++ FeatureExtractor 输出与 Python librosa 输出
"""
import argparse
import numpy as np
import torch
import torchaudio
import torchaudio.functional as F
import librosa


def extract_python_mel(wav_path, sr=16000, n_fft=2048, hop_length=512, n_mels=64, fmax=8000):
    y, sr_loaded = 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)
    # Normalize per mel
    mel_db = (mel_db - mel_db.mean(axis=1, keepdims=True)) / (mel_db.std(axis=1, keepdims=True) + 1e-8)
    # Pad/truncate to 63 frames
    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


def load_cpp_mel(bin_path, n_mels=64, n_frames=63):
    with open(bin_path, 'rb') as f:
        data = np.fromfile(f, dtype=np.float32)
    return data.reshape(n_mels, n_frames)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--wav', required=True)
    parser.add_argument('--cpp_bin', required=True, help='C++ extracted mel binary')
    parser.add_argument('--sr', type=int, default=16000)
    args = parser.parse_args()

    py_mel = extract_python_mel(args.wav, sr=args.sr)
    cpp_mel = load_cpp_mel(args.cpp_bin)

    diff = np.abs(py_mel - cpp_mel)
    print(f"Python shape: {py_mel.shape}, C++ shape: {cpp_mel.shape}")
    print(f"Max absolute diff: {diff.max():.6f}")
    print(f"Mean absolute diff: {diff.mean():.6f}")
    print(f"RMSE: {np.sqrt(np.mean((py_mel - cpp_mel)**2)):.6f}")

    if diff.max() < 0.01:
        print("[PASS] Feature consistency within tolerance.")
    else:
        print("[FAIL] Feature mismatch detected.")


if __name__ == '__main__':
    main()
