#!/usr/bin/env python3
"""
Precompute entropy and perplexity for test-10 audio samples.

This script loads the entropy model once and processes all 10 audio files,
saving tokens, entropies, and perplexities to JSON files for the Gradio app.

Usage:
    uv run python apps/patch_view/scripts/precompute_patch_view.py
"""

import sys
import torch
import numpy as np
import json
import math
from pathlib import Path
import torchaudio

# Add project root and WavTokenizer to path
# __file__ is in apps/patch_view/scripts/, so climb three levels to project root
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "WavTokenizer"))

from encoder.utils import convert_audio
from decoder.pretrained import WavTokenizer
from bytelatent.transformer import LMTransformer, LMTransformerArgs
from torch.distributed.checkpoint.state_dict import get_model_state_dict, set_model_state_dict
import torch.distributed.checkpoint as dcp


def load_entropy_model(checkpoint_dir, device):
    """Load entropy model from checkpoint (step 50000)."""
    print(f"Loading entropy model from {checkpoint_dir}...")

    params_path = Path(checkpoint_dir) / "params.json"
    with open(params_path) as f:
        params = json.load(f)

    model_args = LMTransformerArgs(**params["entropy_model"])
    model = LMTransformer(model_args)

    # Load checkpoint
    state_dict = {"model": get_model_state_dict(model)}
    dcp.load(state_dict, checkpoint_id=checkpoint_dir)
    set_model_state_dict(model, state_dict["model"])

    model = model.to(device)
    model.eval()
    num_params = sum(p.numel() for p in model.parameters())
    print(f"  ✓ Entropy model loaded ({num_params:,} parameters)")

    return model


def load_wavtokenizer(config_path, model_path, device):
    """Load WavTokenizer for audio encoding."""
    wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path)
    wavtokenizer = wavtokenizer.to(device)
    wavtokenizer.eval()
    return wavtokenizer


def encode_audio(audio_path, wavtokenizer, device):
    """
    Load and encode audio file to tokens.

    Returns:
        tokens: List of token IDs
        duration: Duration in seconds
        original_wav: Original audio (numpy array)
        reconstructed_wav: WavTokenizer reconstructed audio (numpy array)
    """
    # Load audio using torchaudio
    wav, sr = torchaudio.load(audio_path)

    # Convert to 24kHz mono
    wav = convert_audio(wav, sr, 24000, 1)

    # Normalize
    wav = wav / (torch.abs(wav).max() + 1e-8)

    # Keep original for comparison
    original_wav = wav.squeeze().cpu().numpy()

    # Encode to tokens
    bandwidth_id = torch.tensor([0]).to(device)
    with torch.no_grad():
        _, discrete_codes = wavtokenizer.encode_infer(wav.to(device), bandwidth_id=bandwidth_id)

    # discrete_codes: [n_q, B, seq_len] - extract first codebook, first batch
    tokens = discrete_codes[0, 0].cpu()

    # Decode tokens to reconstruct audio
    with torch.no_grad():
        tokens_for_decode = tokens.unsqueeze(0).unsqueeze(0).to(device)  # [1, 1, seq_len]
        features = wavtokenizer.codes_to_features(tokens_for_decode)
        reconstructed = wavtokenizer.decode(features, bandwidth_id=bandwidth_id)
        reconstructed_wav = reconstructed.squeeze().cpu().numpy()

    tokens_list = tokens.tolist()
    duration = len(tokens_list) / 40.0  # 40 tokens/sec

    return tokens_list, duration, original_wav, reconstructed_wav


def compute_entropies(tokens, entropy_model, device, max_context=512):
    """
    Compute entropy for each token position using batched contexts.

    Args:
        tokens: List of token IDs
        entropy_model: Loaded entropy model
        max_context: Maximum context window used per forward pass

    Returns:
        numpy array of entropy values (bits) aligned with token positions
    """
    num_tokens = len(tokens)
    print(f"  Computing entropies for {num_tokens} tokens...")

    if num_tokens == 0:
        return np.array([])

    token_tensor = torch.tensor(tokens, dtype=torch.long, device=device)
    entropies = torch.zeros(num_tokens, dtype=torch.float32, device=device)

    chunk_size = max(1, max_context // 2)

    with torch.no_grad():
        processed = 0
        while processed < num_tokens:
            chunk_end = min(num_tokens, processed + chunk_size)
            context_start = max(0, chunk_end - max_context)
            context_tokens = token_tensor[context_start:chunk_end].unsqueeze(0)

            logits = entropy_model(context_tokens)
            probs = torch.softmax(logits[0], dim=-1)
            log_probs = torch.log(probs + 1e-10) / math.log(2)
            chunk_entropy = -(probs * log_probs).sum(dim=-1)

            new_count = chunk_end - processed
            entropies[processed:chunk_end] = chunk_entropy[-new_count:]

            processed = chunk_end
            print(f"    Processed {processed}/{num_tokens} tokens", end="\r")

    print()  # newline after progress carriage-return
    return entropies.cpu().numpy()


def process_audio_file(audio_path, entropy_model, wavtokenizer, device, output_dir):
    """
    Process single audio file: encode, compute entropies, save JSON and reconstructed audio.

    Args:
        audio_path: Path to FLAC file
        entropy_model: Loaded entropy model
        wavtokenizer: Loaded WavTokenizer
        device: torch device
        output_dir: Directory to save JSON

    Returns:
        Path to saved JSON file
    """
    import soundfile as sf

    file_id = audio_path.stem
    print(f"\nProcessing {file_id}...")

    # Load transcript if available
    txt_path = audio_path.with_suffix(".txt")
    transcript = txt_path.read_text().strip() if txt_path.exists() else ""

    # Encode audio and get reconstruction
    print("  Encoding audio...")
    tokens, duration, original_wav, reconstructed_wav = encode_audio(
        audio_path, wavtokenizer, device
    )
    print(f"  ✓ Encoded {len(tokens)} tokens ({duration:.2f}s)")

    # Save reconstructed audio for comparison
    reconstructed_path = output_dir / f"{file_id}_reconstructed.wav"
    sf.write(reconstructed_path, reconstructed_wav, 24000)
    print(f"  ✓ Saved reconstructed audio ({reconstructed_path.stat().st_size / 1024:.1f} KB)")

    # Compute entropies
    entropies = compute_entropies(tokens, entropy_model, device)
    print(f"  ✓ Computed entropies")

    # Compute perplexities: PPL = 2^H
    perplexities = 2 ** entropies
    print(f"  ✓ Computed perplexities")

    # Create output data
    data = {
        "file_id": file_id,
        "audio_path": str(audio_path),
        "reconstructed_audio_path": str(reconstructed_path),
        "transcript": transcript,
        "audio_duration_sec": float(duration),
        "sample_rate": 24000,
        "tokens": tokens,
        "entropies": entropies.tolist(),
        "perplexities": perplexities.tolist(),
        "num_tokens": len(tokens),
        "avg_entropy": float(np.mean(entropies)),
        "min_entropy": float(np.min(entropies)),
        "max_entropy": float(np.max(entropies)),
        "avg_perplexity": float(np.mean(perplexities)),
    }

    # Save to JSON
    output_file = output_dir / f"{file_id}.json"
    with open(output_file, 'w') as f:
        json.dump(data, f, indent=2)

    print(f"  ✓ Saved to {output_file} ({output_file.stat().st_size / 1024:.1f} KB)")

    return output_file


def main():
    # Configuration
    ENTROPY_CHECKPOINT = "archive/checkpoints/audio_peoplespeech_50k_scratch/checkpoints/0000050000"
    WAVTOKENIZER_CONFIG = "WavTokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
    WAVTOKENIZER_MODEL = "WavTokenizer/WavTokenizer_small_600_24k_4096.ckpt"
    TEST_AUDIO_DIR = Path("data/test-10")
    OUTPUT_DIR = Path("data/test-10/precomputed")

    print("=" * 70)
    print("BLT Patch Visualization - Precomputation")
    print("=" * 70)
    print(f"\nEntropy Model: {ENTROPY_CHECKPOINT}")
    print(f"Audio Files: {TEST_AUDIO_DIR}")
    print(f"Output: {OUTPUT_DIR}")

    # Create output directory
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Setup device
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"\nUsing device: {device}")

    # Load models
    print("\n" + "=" * 70)
    print("Loading Models")
    print("=" * 70)

    print("\nLoading WavTokenizer...")
    wavtokenizer = load_wavtokenizer(WAVTOKENIZER_CONFIG, WAVTOKENIZER_MODEL, device)
    print("✓ WavTokenizer loaded")

    print("\nLoading Entropy Model...")
    entropy_model = load_entropy_model(ENTROPY_CHECKPOINT, device)
    print("✓ Entropy Model loaded")

    # Process all audio files
    print("\n" + "=" * 70)
    print("Processing Audio Files")
    print("=" * 70)

    audio_files = sorted(TEST_AUDIO_DIR.glob("*.flac"))
    print(f"\nFound {len(audio_files)} audio files")

    results = []
    for i, audio_path in enumerate(audio_files, 1):
        print(f"\n[{i}/{len(audio_files)}] {audio_path.name}")
        try:
            output_file = process_audio_file(
                audio_path, entropy_model, wavtokenizer, device, OUTPUT_DIR
            )
            results.append({
                "file_id": audio_path.stem,
                "status": "success",
                "output_file": str(output_file)
            })
        except Exception as e:
            print(f"  ✗ Error: {e}")
            results.append({
                "file_id": audio_path.stem,
                "status": "error",
                "error": str(e)
            })

    # Summary
    print("\n" + "=" * 70)
    print("Summary")
    print("=" * 70)

    successful = sum(1 for r in results if r["status"] == "success")
    failed = sum(1 for r in results if r["status"] == "error")

    print(f"\nProcessed: {len(results)} files")
    print(f"  ✓ Successful: {successful}")
    if failed > 0:
        print(f"  ✗ Failed: {failed}")

    # Calculate total size
    total_size = sum(
        (OUTPUT_DIR / f"{r['file_id']}.json").stat().st_size
        for r in results if r["status"] == "success"
    )
    print(f"\nTotal precomputed data size: {total_size / 1024:.1f} KB")
    print(f"Average per file: {total_size / successful / 1024:.1f} KB")

    print(f"\n✓ Precomputation complete!")
    print(f"  Data saved to: {OUTPUT_DIR}/")
    print(f"\nYou can now run the Gradio app:")
    print(f"  uv run python -m patch_view.app")


if __name__ == "__main__":
    main()