#!/usr/bin/env python3
"""
Dataset Validator - Validates the generated dataset for training.

Checks:
1. Structure: list of dicts with required fields
2. Whisper features: shape [seq_len, 1280]
3. SNAC tokens: multiple of 7, valid range
4. Optional fields: text, answer, text_tokens, word_alignments
5. Statistics and distribution
"""

import argparse
import sys
from pathlib import Path
from collections import defaultdict

import torch
import numpy as np


# SNAC token constants (from CLAUDE.md)
SNAC_BASE = 128266
SNAC_LAYERS = 7
SNAC_VOCAB_PER_LAYER = 4096
WHISPER_DIM = 1280


def validate_sample(idx: int, sample: dict, verbose: bool = False) -> tuple[bool, list[str]]:
    """Validate a single sample. Returns (is_valid, list of errors)."""
    errors = []
    warnings = []

    # Check required fields
    if "whisper_features" not in sample:
        errors.append("Missing 'whisper_features'")
    if "snac_tokens" not in sample:
        errors.append("Missing 'snac_tokens'")

    if errors:
        return False, errors

    # Validate whisper_features
    wf = sample["whisper_features"]
    if not isinstance(wf, torch.Tensor):
        errors.append(f"whisper_features should be Tensor, got {type(wf).__name__}")
    else:
        if wf.dim() != 2:
            errors.append(f"whisper_features should be 2D [seq_len, 1280], got {wf.dim()}D")
        elif wf.shape[1] != WHISPER_DIM:
            errors.append(f"whisper_features dim should be {WHISPER_DIM}, got {wf.shape[1]}")
        if wf.shape[0] == 0:
            errors.append("whisper_features has 0 length")
        if torch.isnan(wf).any():
            errors.append("whisper_features contains NaN values")
        if torch.isinf(wf).any():
            errors.append("whisper_features contains Inf values")

    # Validate snac_tokens
    st = sample["snac_tokens"]
    if not isinstance(st, torch.Tensor):
        errors.append(f"snac_tokens should be Tensor, got {type(st).__name__}")
    else:
        if st.dim() != 1:
            errors.append(f"snac_tokens should be 1D, got {st.dim()}D")
        if len(st) == 0:
            errors.append("snac_tokens has 0 length")
        elif len(st) % SNAC_LAYERS != 0:
            errors.append(f"snac_tokens length ({len(st)}) not multiple of {SNAC_LAYERS}")

        # Check token range (raw tokens before offset, should be 0-4095)
        if len(st) > 0:
            min_tok = st.min().item()
            max_tok = st.max().item()
            # Tokens could be raw (0-4095) or with offset applied (128266+)
            if max_tok < SNAC_BASE:
                # Raw tokens
                if min_tok < 0 or max_tok >= SNAC_VOCAB_PER_LAYER:
                    warnings.append(f"snac_tokens range [{min_tok}, {max_tok}] outside [0, {SNAC_VOCAB_PER_LAYER-1}]")
            else:
                # Tokens with offset
                expected_max = SNAC_BASE + (SNAC_LAYERS * SNAC_VOCAB_PER_LAYER)
                if min_tok < SNAC_BASE or max_tok >= expected_max:
                    warnings.append(f"snac_tokens with offset range [{min_tok}, {max_tok}] unexpected")

    # Validate optional fields
    if "text" in sample and not isinstance(sample["text"], str):
        warnings.append(f"text should be str, got {type(sample['text']).__name__}")

    if "answer" in sample and not isinstance(sample["answer"], str):
        warnings.append(f"answer should be str, got {type(sample['answer']).__name__}")

    if "text_tokens" in sample:
        tt = sample["text_tokens"]
        if not isinstance(tt, torch.Tensor):
            warnings.append(f"text_tokens should be Tensor, got {type(tt).__name__}")
        elif tt.dim() != 1:
            warnings.append(f"text_tokens should be 1D, got {tt.dim()}D")

    if "word_alignments" in sample:
        wa = sample["word_alignments"]
        if not isinstance(wa, list):
            warnings.append(f"word_alignments should be list, got {type(wa).__name__}")

    if verbose and warnings:
        for w in warnings:
            print(f"  [WARN] Sample {idx}: {w}")

    return len(errors) == 0, errors


def compute_statistics(dataset: list) -> dict:
    """Compute dataset statistics."""
    stats = {
        "total_samples": len(dataset),
        "whisper_lengths": [],
        "snac_lengths": [],
        "snac_frames": [],
        "has_text": 0,
        "has_answer": 0,
        "has_text_tokens": 0,
        "has_word_alignments": 0,
        "text_lengths": [],
        "answer_lengths": [],
    }

    for sample in dataset:
        if "whisper_features" in sample and isinstance(sample["whisper_features"], torch.Tensor):
            stats["whisper_lengths"].append(sample["whisper_features"].shape[0])

        if "snac_tokens" in sample and isinstance(sample["snac_tokens"], torch.Tensor):
            length = len(sample["snac_tokens"])
            stats["snac_lengths"].append(length)
            stats["snac_frames"].append(length // SNAC_LAYERS)

        if "text" in sample:
            stats["has_text"] += 1
            if isinstance(sample["text"], str):
                stats["text_lengths"].append(len(sample["text"]))

        if "answer" in sample:
            stats["has_answer"] += 1
            if isinstance(sample["answer"], str):
                stats["answer_lengths"].append(len(sample["answer"]))

        if "text_tokens" in sample:
            stats["has_text_tokens"] += 1

        if "word_alignments" in sample:
            stats["has_word_alignments"] += 1

    return stats


def print_statistics(stats: dict):
    """Print dataset statistics."""
    print("\n" + "=" * 60)
    print("DATASET STATISTICS")
    print("=" * 60)

    print(f"\nTotal samples: {stats['total_samples']}")

    # Whisper features
    if stats["whisper_lengths"]:
        wl = np.array(stats["whisper_lengths"])
        print(f"\nWhisper features length:")
        print(f"  Min: {wl.min()}, Max: {wl.max()}, Mean: {wl.mean():.1f}, Std: {wl.std():.1f}")

    # SNAC tokens
    if stats["snac_lengths"]:
        sl = np.array(stats["snac_lengths"])
        sf = np.array(stats["snac_frames"])
        print(f"\nSNAC tokens:")
        print(f"  Tokens - Min: {sl.min()}, Max: {sl.max()}, Mean: {sl.mean():.1f}")
        print(f"  Frames - Min: {sf.min()}, Max: {sf.max()}, Mean: {sf.mean():.1f}")

        # Duration estimate (24kHz, 512 samples per frame = 21.3ms per frame)
        duration_sec = sf * 0.0213
        print(f"  Duration - Min: {duration_sec.min():.1f}s, Max: {duration_sec.max():.1f}s, Mean: {duration_sec.mean():.1f}s")

    # Optional fields
    print(f"\nOptional fields present:")
    print(f"  text: {stats['has_text']}/{stats['total_samples']} ({100*stats['has_text']/stats['total_samples']:.1f}%)")
    print(f"  answer: {stats['has_answer']}/{stats['total_samples']} ({100*stats['has_answer']/stats['total_samples']:.1f}%)")
    print(f"  text_tokens: {stats['has_text_tokens']}/{stats['total_samples']} ({100*stats['has_text_tokens']/stats['total_samples']:.1f}%)")
    print(f"  word_alignments: {stats['has_word_alignments']}/{stats['total_samples']} ({100*stats['has_word_alignments']/stats['total_samples']:.1f}%)")

    # Text lengths
    if stats["text_lengths"]:
        tl = np.array(stats["text_lengths"])
        print(f"\nText (question) length (chars):")
        print(f"  Min: {tl.min()}, Max: {tl.max()}, Mean: {tl.mean():.1f}")

    if stats["answer_lengths"]:
        al = np.array(stats["answer_lengths"])
        print(f"\nAnswer length (chars):")
        print(f"  Min: {al.min()}, Max: {al.max()}, Mean: {al.mean():.1f}")


def validate_dataset(path: str, max_samples: int = None, verbose: bool = False) -> bool:
    """Validate the dataset file."""
    print(f"\nValidating: {path}")
    print("=" * 60)

    # Check file exists
    if not Path(path).exists():
        print(f"[ERROR] File not found: {path}")
        return False

    # Load dataset
    print("Loading dataset...")
    try:
        dataset = torch.load(path, map_location="cpu", weights_only=False)
    except Exception as e:
        print(f"[ERROR] Failed to load dataset: {e}")
        return False

    # Check type
    if not isinstance(dataset, list):
        print(f"[ERROR] Dataset should be a list, got {type(dataset).__name__}")
        return False

    total = len(dataset)
    print(f"Loaded {total} samples")

    if total == 0:
        print("[ERROR] Dataset is empty")
        return False

    # Validate samples
    if max_samples and max_samples < total:
        print(f"Validating first {max_samples} samples...")
        samples_to_check = dataset[:max_samples]
    else:
        print(f"Validating all {total} samples...")
        samples_to_check = dataset

    valid_count = 0
    error_counts = defaultdict(int)

    for idx, sample in enumerate(samples_to_check):
        if not isinstance(sample, dict):
            print(f"[ERROR] Sample {idx}: should be dict, got {type(sample).__name__}")
            error_counts["not_dict"] += 1
            continue

        is_valid, errors = validate_sample(idx, sample, verbose=verbose)
        if is_valid:
            valid_count += 1
        else:
            for err in errors:
                error_counts[err] += 1
                if verbose:
                    print(f"[ERROR] Sample {idx}: {err}")

    # Summary
    checked = len(samples_to_check)
    invalid = checked - valid_count

    print(f"\n{'=' * 60}")
    print("VALIDATION SUMMARY")
    print("=" * 60)
    print(f"Samples checked: {checked}/{total}")
    print(f"Valid: {valid_count} ({100*valid_count/checked:.1f}%)")
    print(f"Invalid: {invalid} ({100*invalid/checked:.1f}%)")

    if error_counts:
        print(f"\nError breakdown:")
        for err, count in sorted(error_counts.items(), key=lambda x: -x[1]):
            print(f"  {count:5d}x {err}")

    # Compute and print statistics
    stats = compute_statistics(samples_to_check)
    print_statistics(stats)

    # Check a sample for inspection
    if verbose and valid_count > 0:
        print(f"\n{'=' * 60}")
        print("SAMPLE INSPECTION (first valid sample)")
        print("=" * 60)
        for idx, sample in enumerate(samples_to_check):
            is_valid, _ = validate_sample(idx, sample)
            if is_valid:
                print(f"Sample {idx}:")
                for key, value in sample.items():
                    if isinstance(value, torch.Tensor):
                        print(f"  {key}: Tensor {value.shape} {value.dtype}")
                    elif isinstance(value, str):
                        preview = value[:100] + "..." if len(value) > 100 else value
                        print(f"  {key}: '{preview}'")
                    elif isinstance(value, list):
                        print(f"  {key}: list[{len(value)}]")
                    else:
                        print(f"  {key}: {type(value).__name__}")
                break

    print("\n" + "=" * 60)
    if invalid == 0:
        print("RESULT: PASSED - All samples valid")
        return True
    else:
        print(f"RESULT: FAILED - {invalid} invalid samples")
        return False


def main():
    parser = argparse.ArgumentParser(description="Validate dataset for training")
    parser.add_argument("--path", type=str, required=True, help="Path to dataset .pt file")
    parser.add_argument("--max-samples", type=int, default=None, help="Max samples to validate (default: all)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    args = parser.parse_args()

    success = validate_dataset(args.path, args.max_samples, args.verbose)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()