#!/usr/bin/env python3
"""
Test semantic encoding on a single audio file.

This script encodes a single audio file to verify the pipeline works correctly.
"""

import numpy as np
import os
import torchaudio

# Import encoding functions
from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models_,
    encode as encode_semantic,
)

# Import Audio class
from suno_utils.audio import Audio


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Test semantic encoding on a single file"
    )
    parser.add_argument(
        "--n_codebooks",
        type=int,
        default=None,
        help="Number of codebooks to use (default: use all available)",
    )
    args = parser.parse_args()

    print("=" * 80)
    print("SINGLE FILE SEMANTIC ENCODING TEST")
    print("=" * 80)
    print()

    # Test file
    test_file = "/app2/suno/data/raw_audio_opus_v0/muFDKWH1aaY.opus"
    test_id = "muFDKWH1aaY"
    output_dir = "/home/tony/Work/tony/RealGen/test_output"
    output_file = os.path.join(output_dir, f"{test_id}.npz")

    print(f"Test file:   {test_file}")
    print(f"Output:      {output_file}")
    print(f"N codebooks: {args.n_codebooks if args.n_codebooks else 'all available'}")
    print()

    # Create output directory
    os.makedirs(output_dir, exist_ok=True)

    # Load semantic model
    print("Loading MERT-25 semantic model...")

    # IMPORTANT: Clear any cached models first to ensure we load with correct centroids
    # (suno_utils caches models by device only, not by centroids filepath)
    from suno_utils.tasks.mert_25 import clean_models

    clean_models()

    preload_semantic_models_(
        checkpoint_filepath="/app/suno/data/dpo/models/mert_25.pt",
        centroids_filepath="/app/suno/data/dpo/models/mert_25_2x4k.npy",
        # centroids_filepath="/home/minz/temp/mert_768d_centroids_4000_50.npy",
        device="cuda:0",
    )
    print("Model loaded!\n")

    # Load audio
    print("Loading audio file...")
    if not os.path.exists(test_file):
        print(f"Error: File not found: {test_file}")
        return

    # Load with torchaudio
    waveform, sample_rate = torchaudio.load(test_file)
    print(
        f"  Original: {sample_rate}Hz, {waveform.shape[0]} channels, {waveform.shape[1]/sample_rate:.2f}s"
    )

    # Convert to mono by averaging channels
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    # Resample to 24kHz
    if sample_rate != 24000:
        waveform = torchaudio.functional.resample(waveform, sample_rate, 24000)

    # Convert to Audio object (mono, 24kHz)
    audio = Audio.from_array_float(
        waveform.numpy().squeeze(), sample_rate=24000, max_allowed_val=100
    )
    print(f"  Converted: 24000Hz, 1 channel, {audio.duration_s:.2f}s")
    print()

    # Encode semantic codes
    print("Encoding semantic codes...")
    codes = encode_semantic(audio, device="cuda:0", n_codebooks=args.n_codebooks)
    print(f"  Encoded shape: {codes.shape}")
    print(f"  Data type: {codes.dtype}")
    print(f"  Value range: [{codes.min()}, {codes.max()}]")
    if codes.ndim == 2:
        print(f"  N codebooks used: {codes.shape[1]}")
    print()

    # Convert to int64 and save
    codes = codes.astype(np.int64)
    np.savez(output_file, codes=codes)
    print(f"Saved to: {output_file}")
    print()

    # Verify saved file
    print("Verifying saved file...")
    loaded = np.load(output_file)
    loaded_codes = loaded["codes"]
    print(f"  Loaded shape: {loaded_codes.shape}")
    print(f"  Loaded dtype: {loaded_codes.dtype}")
    print(f"  Values match: {np.array_equal(codes, loaded_codes)}")
    print()

    # Show sample codes
    print("Sample codes (first 20):")
    print(loaded_codes[:20])
    print()

    print("Sample codes (last 20):")
    print(loaded_codes[-20:])
    print()

    print("=" * 80)
    print("✓ Test completed successfully!")
    print("=" * 80)


if __name__ == "__main__":
    main()