#!/usr/bin/env python3
"""
Verify that enable_efficient_vision_embed produces identical results.

This script tests that the new efficient vision embedding logic produces
the same noisy_embeds as the original logic.

Usage:
    python chengyue_script/dvlm/utils/verify_efficient_vision_embed.py
"""

import torch
import torch.nn as nn


def test_efficient_vision_embed_equivalence():
    """Test that old and new logic produce identical embeddings."""
    
    # Setup
    batch_size = 2
    seq_len = 32
    hidden_size = 128
    vocab_size = 1000
    mask_id = 999
    
    # Create mock embedding layer
    embed_tokens = nn.Embedding(vocab_size, hidden_size)
    
    # Create mock data
    torch.manual_seed(42)
    input_ids = torch.randint(0, vocab_size - 1, (batch_size, seq_len))
    labels = torch.full((batch_size, seq_len), -100, dtype=torch.long)
    # Set some positions as text tokens (labels != -100)
    labels[:, 10:25] = input_ids[:, 10:25]  # Response tokens
    
    # Mark some positions as vision tokens
    vision_token_id = 500
    input_ids[:, 2:8] = vision_token_id  # Vision tokens
    
    # Create original_embeds (simulating vision encoder output)
    original_embeds = embed_tokens(input_ids)
    # Simulate vision embeddings being different from text embeddings
    vision_token_mask = (input_ids == vision_token_id)
    vision_mask_3d = vision_token_mask.unsqueeze(-1).expand(-1, -1, hidden_size)
    original_embeds[vision_token_mask] = torch.randn(vision_token_mask.sum(), hidden_size)
    
    # Create x_t (noisy token ids)
    mask_indices = torch.rand(batch_size, seq_len) < 0.5
    x_t = torch.where(mask_indices, torch.full_like(input_ids, mask_id), input_ids)
    
    # Create noisy_input_ids (original logic)
    noisy_input_ids = input_ids.clone()
    noisy_input_ids[labels != -100] = x_t[labels != -100]
    
    # === Original Logic ===
    noisy_embeds_raw = embed_tokens(noisy_input_ids)
    noisy_embeds_old = torch.where(vision_mask_3d, original_embeds, noisy_embeds_raw)
    
    # === New Logic (enable_efficient_vision_embed=True) ===
    noisy_embeds_new = original_embeds.clone()
    text_mask_3d = (labels != -100).unsqueeze(-1).expand(-1, -1, hidden_size)
    noisy_embeds_new = torch.where(text_mask_3d, embed_tokens(x_t), noisy_embeds_new)
    
    # === Verification ===
    print("=" * 60)
    print("Testing enable_efficient_vision_embed equivalence")
    print("=" * 60)
    
    # Check vision token positions
    vision_old = noisy_embeds_old[vision_token_mask]
    vision_new = noisy_embeds_new[vision_token_mask]
    vision_original = original_embeds[vision_token_mask]
    
    vision_match = torch.allclose(vision_old, vision_new, atol=1e-6)
    vision_preserved = torch.allclose(vision_old, vision_original, atol=1e-6)
    print(f"Vision token embeddings match: {vision_match}")
    print(f"Vision embeddings preserved from original: {vision_preserved}")
    
    # Check prompt positions (labels == -100, not vision)
    prompt_mask = (labels == -100) & (~vision_token_mask)
    if prompt_mask.any():
        prompt_old = noisy_embeds_old[prompt_mask]
        prompt_new = noisy_embeds_new[prompt_mask]
        prompt_original = original_embeds[prompt_mask]
        
        prompt_match = torch.allclose(prompt_old, prompt_new, atol=1e-6)
        prompt_preserved = torch.allclose(prompt_old, prompt_original, atol=1e-6)
        print(f"Prompt text embeddings match: {prompt_match}")
        print(f"Prompt embeddings preserved from original: {prompt_preserved}")
    
    # Check text token positions (labels != -100)
    text_mask = (labels != -100)
    text_old = noisy_embeds_old[text_mask]
    text_new = noisy_embeds_new[text_mask]
    
    text_match = torch.allclose(text_old, text_new, atol=1e-6)
    print(f"Text token (response) embeddings match: {text_match}")
    
    # Overall check
    overall_match = torch.allclose(noisy_embeds_old, noisy_embeds_new, atol=1e-6)
    print("-" * 60)
    print(f"OVERALL RESULT: {'PASS ✅' if overall_match else 'FAIL ❌'}")
    print("-" * 60)
    
    if not overall_match:
        diff = (noisy_embeds_old - noisy_embeds_new).abs()
        print(f"Max difference: {diff.max().item():.2e}")
        print(f"Mean difference: {diff.mean().item():.2e}")
        
        # Find where they differ
        diff_positions = diff.sum(-1) > 1e-6
        print(f"Positions with difference: {diff_positions.sum().item()}")
        
    return overall_match


def test_with_various_configs():
    """Test with various configurations."""
    print("\n" + "=" * 60)
    print("Testing with various configurations")
    print("=" * 60)
    
    configs = [
        {"name": "Basic", "mask_ratio": 0.5},
        {"name": "High mask ratio", "mask_ratio": 0.9},
        {"name": "Low mask ratio", "mask_ratio": 0.1},
        {"name": "All masked", "mask_ratio": 1.0},
        {"name": "None masked", "mask_ratio": 0.0},
    ]
    
    all_passed = True
    for config in configs:
        print(f"\nConfig: {config['name']} (mask_ratio={config['mask_ratio']})")
        passed = test_efficient_vision_embed_equivalence()
        all_passed = all_passed and passed
    
    print("\n" + "=" * 60)
    print(f"ALL TESTS: {'PASSED ✅' if all_passed else 'FAILED ❌'}")
    print("=" * 60)
    
    return all_passed


if __name__ == "__main__":
    success = test_efficient_vision_embed_equivalence()
    exit(0 if success else 1)

