#!/usr/bin/env python3
"""
Verify dataset quality by transcribing audio and comparing with original text.
Uses WhisperX to transcribe TTS-generated audio.
"""
import os
import sys
import torch
import numpy as np
from difflib import SequenceMatcher

# Patch torch.load
_orig = torch.load
torch.load = lambda *a, **kw: _orig(*a, **{**kw, 'weights_only': False})


def similarity(a: str, b: str) -> float:
    """Calculate similarity ratio between two strings."""
    return SequenceMatcher(None, a.lower(), b.lower()).ratio()


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Verify dataset audio quality")
    parser.add_argument("--dataset", type=str, default="./data/test_b200.pt", help="Dataset path")
    parser.add_argument("--samples", type=int, default=10, help="Number of samples to verify")
    parser.add_argument("--gpu", type=int, default=0, help="GPU to use")
    args = parser.parse_args()

    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

    import torchaudio
    import whisperx

    print("=" * 60)
    print("Dataset Audio Verification")
    print("=" * 60)

    # Load dataset
    print(f"\nLoading dataset: {args.dataset}")
    data = torch.load(args.dataset)
    print(f"Total items: {len(data)}")

    # Load TTS for regenerating audio
    print("\nLoading TTS...")
    from soprano import SopranoTTS
    tts = SopranoTTS(backend="transformers", device="cuda")

    # Load WhisperX for transcription
    print("Loading WhisperX...")
    wx_model = whisperx.load_model("large-v3-turbo", "cuda", compute_type="float16", language="en")

    # Select random samples
    import random
    indices = random.sample(range(len(data)), min(args.samples, len(data)))

    print(f"\nVerifying {len(indices)} samples...\n")
    print("-" * 60)

    results = {
        "question": {"total": 0, "good": 0, "similarities": []},
        "answer": {"total": 0, "good": 0, "similarities": []}
    }

    for i, idx in enumerate(indices):
        item = data[idx]
        q_text = item["text"]
        a_text = item["answer"]

        print(f"[{i+1}/{len(indices)}] Sample {idx}")

        # Generate audio with TTS
        q_audio = tts.infer(q_text)
        a_audio = tts.infer(a_text)

        # Convert to numpy
        q_np = q_audio.cpu().numpy() if hasattr(q_audio, 'cpu') else np.array(q_audio)
        a_np = a_audio.cpu().numpy() if hasattr(a_audio, 'cpu') else np.array(a_audio)

        # Resample to 16kHz for WhisperX
        q_16k = torchaudio.functional.resample(torch.from_numpy(q_np), 32000, 16000).numpy()
        a_16k = torchaudio.functional.resample(torch.from_numpy(a_np), 32000, 16000).numpy()

        # Transcribe
        q_result = wx_model.transcribe(q_16k, batch_size=1)
        a_result = wx_model.transcribe(a_16k, batch_size=1)

        q_transcribed = " ".join([s["text"].strip() for s in q_result["segments"]])
        a_transcribed = " ".join([s["text"].strip() for s in a_result["segments"]])

        # Calculate similarity
        q_sim = similarity(q_text, q_transcribed)
        a_sim = similarity(a_text, a_transcribed)

        results["question"]["total"] += 1
        results["answer"]["total"] += 1
        results["question"]["similarities"].append(q_sim)
        results["answer"]["similarities"].append(a_sim)

        if q_sim >= 0.8:
            results["question"]["good"] += 1
        if a_sim >= 0.8:
            results["answer"]["good"] += 1

        # Print results
        q_status = "✓" if q_sim >= 0.8 else "✗"
        a_status = "✓" if a_sim >= 0.8 else "✗"

        print(f"  Question ({q_sim:.0%}) {q_status}")
        print(f"    Original:    \"{q_text[:60]}...\"" if len(q_text) > 60 else f"    Original:    \"{q_text}\"")
        print(f"    Transcribed: \"{q_transcribed[:60]}...\"" if len(q_transcribed) > 60 else f"    Transcribed: \"{q_transcribed}\"")

        print(f"  Answer ({a_sim:.0%}) {a_status}")
        print(f"    Original:    \"{a_text[:60]}...\"" if len(a_text) > 60 else f"    Original:    \"{a_text}\"")
        print(f"    Transcribed: \"{a_transcribed[:60]}...\"" if len(a_transcribed) > 60 else f"    Transcribed: \"{a_transcribed}\"")
        print()

    # Summary
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)

    q_avg = np.mean(results["question"]["similarities"])
    a_avg = np.mean(results["answer"]["similarities"])

    print(f"\nQuestions:")
    print(f"  Good (>=80%): {results['question']['good']}/{results['question']['total']}")
    print(f"  Avg similarity: {q_avg:.1%}")

    print(f"\nAnswers:")
    print(f"  Good (>=80%): {results['answer']['good']}/{results['answer']['total']}")
    print(f"  Avg similarity: {a_avg:.1%}")

    overall = (q_avg + a_avg) / 2
    print(f"\nOverall similarity: {overall:.1%}")

    if overall >= 0.85:
        print("\n✓ Dataset quality: GOOD")
    elif overall >= 0.70:
        print("\n⚠ Dataset quality: ACCEPTABLE")
    else:
        print("\n✗ Dataset quality: POOR")


if __name__ == "__main__":
    main()