#!/usr/bin/env python3
"""
End-to-end test comparing PyTorch SAM Audio with ONNX Runtime.

This script:
1. Loads a real audio sample from AudioCaps
2. Runs PyTorch inference using the original SAMAudio model
3. Runs ONNX inference using the exported models
4. Compares the output waveforms
"""

import torch
import torchaudio
import numpy as np
import os
from datasets import load_dataset


def load_audiocaps_sample():
    """Load a sample from AudioCaps dataset."""
    print("Loading AudioCaps sample...")
    dset = load_dataset(
        "parquet",
        data_files="hf://datasets/OpenSound/AudioCaps/data/test-00000-of-00041.parquet",
    )
    sample = dset["train"][8]["audio"].get_all_samples()
    print(f"  Sample rate: {sample.sample_rate}")
    print(f"  Duration: {sample.data.shape[-1] / sample.sample_rate:.2f}s")
    return sample


def run_pytorch_inference(sample, device="cpu"):
    """Run inference using PyTorch SAMAudio model."""
    print("\n=== PyTorch Inference ===")
    
    from sam_audio import SAMAudio, SAMAudioProcessor
    
    # Load model and processor
    print("Loading SAMAudio model...")
    model = SAMAudio.from_pretrained("facebook/sam-audio-small").to(device).eval()
    processor = SAMAudioProcessor.from_pretrained("facebook/sam-audio-small")
    
    # Resample and prepare input
    wav = torchaudio.functional.resample(
        sample.data, sample.sample_rate, processor.audio_sampling_rate
    )
    wav = wav.mean(0, keepdim=True)  # Convert to mono
    
    print(f"  Input audio shape: {wav.shape}")
    print(f"  Sample rate: {processor.audio_sampling_rate}")
    
    # Prepare inputs with explicit anchor
    inputs = processor(
        audios=[wav],
        descriptions=["A horn honking"],
        anchors=[[["+", 6.3, 7.0]]]
    ).to(device)
    
    # Run separation
    print("Running separation...")
    with torch.inference_mode():
        result = model.separate(inputs)
    
    separated_audio = result.target[0].cpu().numpy()
    print(f"  Output shape: {separated_audio.shape}")
    
    return separated_audio, processor.audio_sampling_rate, wav.numpy()


def run_onnx_inference(sample, model_dir="."):
    """Run inference using ONNX models."""
    print("\n=== ONNX Runtime Inference ===")
    
    import onnxruntime as ort
    from transformers import AutoTokenizer
    import json
    
    # Load models
    print("Loading ONNX models...")
    providers = ["CPUExecutionProvider"]
    
    dacvae_encoder = ort.InferenceSession(
        os.path.join(model_dir, "dacvae_encoder.onnx"),
        providers=providers,
    )
    dacvae_decoder = ort.InferenceSession(
        os.path.join(model_dir, "dacvae_decoder.onnx"),
        providers=providers,
    )
    t5_encoder = ort.InferenceSession(
        os.path.join(model_dir, "t5_encoder.onnx"),
        providers=providers,
    )
    dit = ort.InferenceSession(
        os.path.join(model_dir, "dit_single_step.onnx"),
        providers=providers,
    )
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_dir, "tokenizer"))
    print("  All models loaded")
    
    # Prepare audio (resample to 44.1kHz for DACVAE)
    wav = torchaudio.functional.resample(
        sample.data, sample.sample_rate, 44100
    )
    wav = wav.mean(0, keepdim=True)  # Convert to mono
    audio = wav.numpy().reshape(1, 1, -1).astype(np.float32)
    
    print(f"  Input audio shape: {audio.shape}")
    
    # 1. Encode audio
    print("Encoding audio...")
    latent = dacvae_encoder.run(
        ["latent_features"],
        {"audio": audio}
    )[0]
    print(f"  Audio latent shape: {latent.shape}")
    
    # 2. Encode text
    print("Encoding text...")
    tokens = tokenizer(
        "A horn honking",
        return_tensors="np",
        padding=True,
        truncation=True,
        max_length=77,
    )
    text_features = t5_encoder.run(
        ["hidden_states"],
        {
            "input_ids": tokens["input_ids"].astype(np.int64),
            "attention_mask": tokens["attention_mask"].astype(np.int64),
        }
    )[0]
    print(f"  Text features shape: {text_features.shape}")
    
    # 3. Run ODE solving (simplified - just one step for testing)
    print("Running DiT (simplified test - 1 step)...")
    batch_size = 1
    latent_dim = latent.shape[1]  # 128
    time_steps = latent.shape[2]
    
    # Prepare inputs
    # SAMAudio._get_audio_features: returns torch.cat([audio_features, audio_features], dim=2)
    # So audio_features is the mixture DUPLICATED, not mixture + zeros!
    mixture_features = latent.transpose(0, 2, 1)  # (B, T, 128) - from DACVAE
    
    # Duplicate mixture features (this is what SAMAudio actually does)
    audio_features = np.concatenate([
        mixture_features,  # Mixture latent
        mixture_features   # Mixture latent (DUPLICATE - not zeros!)
    ], axis=-1)  # -> (B, T, 256)
    
    # noisy_audio starts from random noise for ODE solving from t=0 to t=1
    # SAMAudio uses: noise = torch.randn_like(audio_features)
    initial = np.random.randn(batch_size, time_steps, 256).astype(np.float32)
    
    # Just run one step to verify the model works
    velocity = dit.run(
        ["velocity"],
        {
            "noisy_audio": initial,
            "time": np.array([0.0], dtype=np.float32),
            "audio_features": audio_features,
            "text_features": text_features,
            "text_mask": tokens["attention_mask"].astype(bool),
            "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
            "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
            "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
            "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
        }
    )[0]
    print(f"  DiT velocity shape: {velocity.shape}")

    
    # 4. Run full ODE solve (16 steps midpoint method)
    print("Running full ODE solve (16 steps)...")
    num_steps = 16
    dt = 1.0 / num_steps
    x = initial.copy()
    
    for i in range(num_steps):
        t = np.array([i * dt], dtype=np.float32)
        t_mid = np.array([t[0] + dt / 2], dtype=np.float32)
        
        # k1 = f(t, x)
        k1 = dit.run(
            ["velocity"],
            {
                "noisy_audio": x,
                "time": t,
                "audio_features": audio_features,
                "text_features": text_features,
                "text_mask": tokens["attention_mask"].astype(bool),
                "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
                "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
                "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
                "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
            }
        )[0]
        
        # Midpoint
        x_mid = x + (dt / 2) * k1
        
        # k2 = f(t_mid, x_mid)
        k2 = dit.run(
            ["velocity"],
            {
                "noisy_audio": x_mid,
                "time": t_mid,
                "audio_features": audio_features,
                "text_features": text_features,
                "text_mask": tokens["attention_mask"].astype(bool),
                "masked_video_features": np.zeros((batch_size, 1024, time_steps), dtype=np.float32),
                "anchor_ids": np.zeros((batch_size, time_steps), dtype=np.int64),
                "anchor_alignment": np.zeros((batch_size, time_steps), dtype=np.int64),
                "audio_pad_mask": np.ones((batch_size, time_steps), dtype=bool),
            }
        )[0]
        
        # Update
        x = x + dt * k2
        print(f"  Step {i+1}/{num_steps}")
    
    # 5. Extract separated latent and decode in chunks
    # (The DACVAE decoder was exported with fixed time=25, so we decode in chunks)
    print("Decoding audio...")
    # SAMAudio: target is first 128 dims, residual is second 128 dims
    # generated_features.reshape(2*B, C, T) -> first B = channels 0:128 (target)
    target_latent = x[:, :, :latent_dim].transpose(0, 2, 1)  # (B, 128, T) - TARGET
    separated_latent = target_latent
    
    # The decoder expects chunks of 25 time steps
    chunk_size = 25
    T = separated_latent.shape[2]
    
    # Process in chunks and concatenate
    audio_chunks = []
    for start_idx in range(0, T, chunk_size):
        end_idx = min(start_idx + chunk_size, T)
        chunk = separated_latent[:, :, start_idx:end_idx]
        
        # Pad last chunk if needed
        actual_size = chunk.shape[2]
        if actual_size < chunk_size:
            pad_size = chunk_size - actual_size
            chunk = np.pad(chunk, ((0, 0), (0, 0), (0, pad_size)), mode='constant')
        
        chunk_audio = dacvae_decoder.run(
            ["waveform"],
            {"latent_features": chunk.astype(np.float32)}
        )[0]
        
        # For padded chunks, trim the output
        if actual_size < chunk_size:
            # Each time step produces hop_length (1920) samples at 48kHz
            samples_per_step = 1920
            trim_samples = actual_size * samples_per_step
            chunk_audio = chunk_audio[:, :, :trim_samples]
        
        audio_chunks.append(chunk_audio)
        print(f"    Decoded chunk {start_idx//chunk_size + 1}/{(T + chunk_size - 1)//chunk_size}")
    
    # Concatenate all chunks
    separated_audio = np.concatenate(audio_chunks, axis=2)
    
    print(f"  Output audio shape: {separated_audio.shape}")
    
    return separated_audio.squeeze(), 44100



def compare_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr):
    """Compare PyTorch and ONNX outputs."""
    print("\n=== Comparison ===")
    
    import scipy.signal
    
    # Resample to same rate if needed
    if pytorch_sr != onnx_sr:
        print(f"Resampling PyTorch output from {pytorch_sr} to {onnx_sr}...")
        # Use scipy for resampling
        num_samples = int(len(pytorch_audio) * onnx_sr / pytorch_sr)
        pytorch_audio_resampled = scipy.signal.resample(pytorch_audio, num_samples)
    else:
        pytorch_audio_resampled = pytorch_audio
    
    # Trim to same length
    min_len = min(len(pytorch_audio_resampled), len(onnx_audio))
    pytorch_trimmed = pytorch_audio_resampled[:min_len]
    onnx_trimmed = onnx_audio[:min_len]
    
    # Compute differences
    diff = np.abs(pytorch_trimmed - onnx_trimmed)
    max_diff = diff.max()
    mean_diff = diff.mean()
    
    # Compute correlation
    correlation = np.corrcoef(pytorch_trimmed, onnx_trimmed)[0, 1]
    
    print(f"  PyTorch audio length: {len(pytorch_audio)} samples")
    print(f"  ONNX audio length: {len(onnx_audio)} samples")
    print(f"  Max difference: {max_diff:.6f}")
    print(f"  Mean difference: {mean_diff:.6f}")
    print(f"  Correlation: {correlation:.6f}")
    
    return max_diff, mean_diff, correlation


def save_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr, input_audio, input_sr):
    """Save audio outputs for listening comparison."""
    import soundfile as sf
    
    output_dir = "test_outputs"
    os.makedirs(output_dir, exist_ok=True)
    
    # Save input
    sf.write(os.path.join(output_dir, "input.wav"), input_audio.squeeze(), input_sr)
    print(f"Saved input to {output_dir}/input.wav")
    
    # Save PyTorch output
    sf.write(os.path.join(output_dir, "pytorch_output.wav"), pytorch_audio, pytorch_sr)
    print(f"Saved PyTorch output to {output_dir}/pytorch_output.wav")
    
    # Save ONNX output
    sf.write(os.path.join(output_dir, "onnx_output.wav"), onnx_audio, onnx_sr)
    print(f"Saved ONNX output to {output_dir}/onnx_output.wav")


def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="End-to-end SAM Audio test")
    parser.add_argument("--model-dir", default=".", help="ONNX model directory")
    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    parser.add_argument("--save-outputs", action="store_true", help="Save audio files")
    parser.add_argument("--skip-pytorch", action="store_true", help="Skip PyTorch inference")
    args = parser.parse_args()
    
    # Load sample
    sample = load_audiocaps_sample()
    
    # Run PyTorch inference
    if not args.skip_pytorch:
        pytorch_audio, pytorch_sr, input_audio = run_pytorch_inference(sample, args.device)
    else:
        print("\nSkipping PyTorch inference")
        pytorch_audio, pytorch_sr = None, None
        input_audio = sample.data.mean(0).numpy()
    
    # Run ONNX inference
    onnx_audio, onnx_sr = run_onnx_inference(sample, args.model_dir)
    
    # Compare outputs
    if pytorch_audio is not None:
        compare_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr)
    
    # Save outputs
    if args.save_outputs:
        print("\n=== Saving Outputs ===")
        if pytorch_audio is not None:
            save_outputs(pytorch_audio, onnx_audio, pytorch_sr, onnx_sr, 
                        input_audio, sample.sample_rate)
        else:
            import soundfile as sf
            os.makedirs("test_outputs", exist_ok=True)
            sf.write("test_outputs/onnx_output.wav", onnx_audio, onnx_sr)
            print("Saved ONNX output to test_outputs/onnx_output.wav")
    
    print("\n✓ End-to-end test complete!")


if __name__ == "__main__":
    main()