#!/usr/bin/env python3
"""
Test to verify that decode(latents) and decode(tokens) are now consistent
after fixing the encode() to return quantized latents.
"""

import numpy as np
import torch

from pat.models.pat_tokenizer import PATTokenizer

def test_decode_consistency():
    """Test that decode from latents and tokens gives same result."""
    
    # Load tokenizer
    tokenizer = PATTokenizer.from_pretrained(
        "output/pat_OAT/checkpoints/checkpoint-000660",
        device="cuda" if torch.cuda.is_available() else "cpu"
    )
    tokenizer.eval()
    
    # Create dummy actions
    actions = np.random.randn(2, 30, 20).astype(np.float32)
    
    print("=" * 80)
    print("Test: Decode Consistency (After Fix)")
    print("=" * 80)
    
    # Encode to get both latents and tokens
    # Note: encode() now returns quantized_latents by default (use_quantizer=True)
    quantized_latents, tokens = tokenizer.encode(actions, return_tokens=True)
    
    print(f"Input actions shape: {actions.shape}")
    print(f"Quantized latents shape: {quantized_latents.shape}")
    print(f"Tokens shape: {tokens.shape}")
    print(f"Tokens range: [{tokens.min()}, {tokens.max()}]")
    
    # Decode from quantized latents
    actions_from_quantized, _ = tokenizer.decode(quantized_latents)
    
    # Decode from tokens (discrete)
    actions_from_tokens, _ = tokenizer.decode(tokens, is_tokens=True)
    
    print("\n" + "=" * 80)
    print("Comparison Results:")
    print("=" * 80)
    
    # Compare actions from quantized latents vs tokens
    mse_quantized_tokens = np.mean((actions_from_quantized - actions_from_tokens) ** 2)
    print(f"\nMSE(decode(quantized_latents), decode(tokens)): {mse_quantized_tokens:.10f}")
    print(f"  This should be ZERO (or very close to zero)")
    print(f"  because quantized_latents and tokens represent the SAME information")
    
    if mse_quantized_tokens < 1e-10:
        print(f"  ✓ PASS: Results are identical!")
    else:
        print(f"  ✗ FAIL: Results differ!")
    
    # Compare reconstruction quality
    mse_orig_quantized = np.mean((actions - actions_from_quantized) ** 2)
    mse_orig_tokens = np.mean((actions - actions_from_tokens) ** 2)
    print(f"\n" + "=" * 80)
    print("Reconstruction Quality:")
    print("=" * 80)
    print(f"MSE(original, from_quantized_latents): {mse_orig_quantized:.8f}")
    print(f"MSE(original, from_tokens): {mse_orig_tokens:.8f}")
    print(f"  Both should be similar (training quality)")
    
    # Verify that quantized_latents == indices_to_embedding(tokens)
    print(f"\n" + "=" * 80)
    print("Verification:")
    print("=" * 80)
    
    # Manually reconstruct latents from tokens
    tokens_tensor = torch.from_numpy(tokens).long().to(tokenizer.device)
    fsq = tokenizer.model.fsq_quantizer
    reconstructed_latents = fsq.indices_to_embedding(tokens_tensor).cpu().numpy()
    
    # Check if quantized_latents and reconstructed_latents are the same
    diff = np.abs(quantized_latents - reconstructed_latents).mean()
    print(f"Mean absolute difference (quantized_latents vs indices_to_embedding(tokens)): {diff:.10f}")
    print(f"  This should be ZERO")
    
    if diff < 1e-10:
        print(f"  ✓ PASS: quantized_latents == indices_to_embedding(tokens)")
    else:
        print(f"  ✗ FAIL: They differ!")
    
    print("\n" + "=" * 80)
    print("Conclusion:")
    print("=" * 80)
    print("After the fix:")
    print("✓ encode() returns quantized_latents (matching original OAT)")
    print("✓ decode(quantized_latents) == decode(tokens)  [CORRECT]")
    print("✓ quantized_latents == indices_to_embedding(tokens)  [CONSISTENT]")
    print("=" * 80)


if __name__ == "__main__":
    test_decode_consistency()
