#!/usr/bin/env python3
"""
Full Attention BitTransformerLM Diffusion Inference Test
========================================================

Test the newly trained full bi-directional attention BitTransformerLM model
using denoising diffusion generation to evaluate improvements from full attention training.

Model Configuration:
- Same full bi-directional unchunked attention as training (chunk_size=None)
- Proper eval() mode with dropout management
- Use latest checkpoint_best.pt from full attention training
- Test with same diffusion inference that worked before
"""

import sys
import torch
import torch.nn.functional as F
from datetime import datetime

sys.path.append('/data')
sys.path.append('/data/BitTransformerLM')

from bit_transformer import (
    BitTransformerLM,
    text_to_bits,
    bits_to_text,
    diffusion_inference,
    set_dropout
)

def load_full_attention_model():
    """Load the newly trained full attention BitTransformerLM model."""
    print("🚀 Loading Full Attention BitTransformerLM for diffusion inference...")
    
    # Create model with SAME configuration as full attention training
    model = BitTransformerLM(
        d_model=512,                    # Same as training
        nhead=16,                       # Same as training
        num_layers=8,                   # Same as training
        dim_feedforward=1024,           # Same as training
        max_seq_len=512,               # Same as training
        reversible=True,                # Same as training
        use_checkpoint=False,           # Disable for inference
        use_autocast=False,            # Disable for inference  
        use_act=True,                  # Same as training
        act_threshold=0.9,             # Same as training
        lambda_K=0.05,                 # Same as training
        lambda_C=0.05,                 # Same as training
        lambda_S=0.05,                 # Same as training
        chunk_size=None,               # FULL ATTENTION - same as training
        overlap=0,                     # Same as training
        full_attn_logging=True         # Same as training
    )
    
    # Load the latest checkpoint_best.pt (should be from full attention training)
    checkpoint_path = '/data/BitTransformerLM/checkpoints/checkpoint_best.pt'
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])
    
    # Set to evaluation mode with proper dropout
    model.eval()
    set_dropout(model, 0.0)  # Disable dropout for inference
    
    # Get checkpoint info
    epoch = checkpoint.get('epoch', 'unknown')
    loss = checkpoint.get('loss', 'unknown')
    
    print(f"✅ Full Attention Model loaded! Epoch: {epoch}, Loss: {loss}")
    
    # Calculate parameters
    total_params = sum(p.numel() for p in model.parameters())
    print(f"📊 Parameters: {total_params:,}")
    
    return model

def test_basic_diffusion_generation(model):
    """Test basic unconditional diffusion generation."""
    print("\n🧪 === BASIC FULL ATTENTION DIFFUSION GENERATION ===")
    
    results = []
    
    test_configs = [
        {"length": 36, "steps": 8, "schedule": "linear"},
        {"length": 45, "steps": 12, "schedule": "cosine"}, 
        {"length": 54, "steps": 16, "schedule": "exp"}
    ]
    
    for i, config in enumerate(test_configs, 1):
        print(f"\n--- Test {i}: {config['length']//9} chars, {config['schedule']} ---")
        
        try:
            # Generate with diffusion
            generated_bits = diffusion_inference(
                model,
                length=config['length'],
                steps=config['steps'],
                batch_size=1,
                schedule=config['schedule']
            )
            
            # Try to decode
            bit_list = generated_bits.squeeze().tolist()
            decoded_text = bits_to_text(bit_list)
            
            print(f"✅ SUCCESS: '{decoded_text}'")
            results.append({
                "test": f"basic_{i}",
                "config": config,
                "success": True,
                "output": decoded_text,
                "bits": len(bit_list)
            })
            
        except Exception as e:
            print(f"❌ FAILED: {e}")
            results.append({
                "test": f"basic_{i}",
                "config": config,
                "success": False,
                "error": str(e)
            })
    
    return results

def test_conditioned_diffusion_generation(model):
    """Test prompt-conditioned diffusion generation."""
    print("\n🎯 === CONDITIONED FULL ATTENTION DIFFUSION GENERATION ===")
    
    results = []
    
    test_prompts = [
        "Hello",
        "Hi there", 
        "What is your name?",
        "The weather is",
        "I am",
        "Yes",
        "No"
    ]
    
    for prompt in test_prompts:
        print(f"\n--- Prompt: '{prompt}' ---")
        
        try:
            # Convert prompt to bits
            prompt_bits = text_to_bits(prompt)
            
            # Generate continuation with diffusion (no init_bits - let it generate freely)
            continuation_length = 45  # 5 character continuation  
            generated_bits = diffusion_inference(
                model,
                length=continuation_length,
                steps=12,
                batch_size=1,
                init_bits=None,
                schedule="cosine"
            )
            
            # Combine prompt + generated continuation
            full_bits = prompt_bits + generated_bits.squeeze().tolist()
            
            # Decode continuation only
            continuation_bits = generated_bits.squeeze().tolist()
            continuation_text = bits_to_text(continuation_bits)
            
            # Show combined result
            combined_text = prompt + continuation_text
            print(f"✅ SUCCESS: '{prompt}' → '{combined_text}'")
            results.append({
                "test": "conditioned",
                "prompt": prompt,
                "success": True,
                "full_output": combined_text,
                "continuation": continuation_text,
                "bits": len(continuation_bits)
            })
            
        except Exception as e:
            print(f"❌ FAILED: {e}")
            results.append({
                "test": "conditioned",
                "prompt": prompt,
                "success": False,
                "error": str(e)
            })
    
    return results

def test_code_diffusion_completion(model):
    """Test code/math completion with diffusion."""
    print("\n💻 === CODE COMPLETION FULL ATTENTION DIFFUSION ===")
    
    results = []
    
    test_cases = [
        # Math equations
        "2 + 2 =",
        "1 + 1 =", 
        "5 * 3 =",
        "10 / 2 =",
        
        # Programming constructs
        "def hello():",
        "if x ==",
        "for i in",
        "print(",
        "return",
        
        # Patterns
        "a, b, c,", 
        "1, 2, 3,",
        "function(",
        "var x =",
    ]
    
    for code in test_cases:
        print(f"\n--- Code: '{code}' ---")
        
        try:
            # Convert to bits
            code_bits = text_to_bits(code)
            
            # Generate completion with diffusion (no init_bits)
            completion_length = 45  # 5 character completion
            generated_bits = diffusion_inference(
                model,
                length=completion_length,
                steps=10,
                batch_size=1,
                init_bits=None,
                schedule="linear"
            )
            
            # Decode completion
            completion_bits = generated_bits.squeeze().tolist()
            completion = bits_to_text(completion_bits)
            
            # Show combined result
            combined_text = code + completion
            print(f"✅ SUCCESS: '{code}' → '{combined_text}'")
            
            # Analyze completion
            analysis = []
            if any(c.isalnum() for c in completion):
                analysis.append("Contains alphanumeric")
                print(f"   📊 Analysis: Contains alphanumeric")
            if any(c in "0123456789" for c in completion):
                analysis.append("Contains numbers")
                print(f"   🔢 Analysis: Contains numbers")
            if any(c in "=(){}[];," for c in completion):
                analysis.append("Contains code symbols")
                print(f"   💻 Analysis: Contains code symbols")
            
            results.append({
                "test": "code_completion",
                "prompt": code,
                "success": True,
                "full_output": combined_text,
                "completion": completion,
                "analysis": analysis,
                "bits": len(completion_bits)
            })
            
        except Exception as e:
            print(f"❌ FAILED: {e}")
            results.append({
                "test": "code_completion",
                "prompt": code,
                "success": False,
                "error": str(e)
            })
    
    return results

def compare_with_previous_results():
    """Note about comparison with previous results."""
    print("\n⚖️  === COMPARISON WITH PREVIOUS RESULTS ===")
    print("Previous chunked attention model achieved:")
    print("- Basic generation: 3/3 success (100%)")
    print("- Conditioned generation: 7/7 success (100%)")  
    print("- Code completion: 13/13 success (100%)")
    print("- All diffusion inference succeeded vs 0% autoregressive")
    print("\nTesting if full attention training improved quality...")

def main():
    print("🚀 FULL ATTENTION BITRANSFORMERLM DIFFUSION INFERENCE TEST")
    print("=" * 70)
    print("Testing newly trained full bi-directional attention model")
    print("with denoising diffusion generation")
    print("=" * 70)
    
    # Load model
    model = load_full_attention_model()
    
    # Run tests
    basic_results = test_basic_diffusion_generation(model)
    conditioned_results = test_conditioned_diffusion_generation(model)
    code_results = test_code_diffusion_completion(model)
    
    # Show comparison
    compare_with_previous_results()
    
    # Calculate summary stats
    total_tests = len(basic_results) + len(conditioned_results) + len(code_results)
    successful_tests = sum(1 for r in basic_results + conditioned_results + code_results if r.get('success', False))
    success_rate = (successful_tests / total_tests) * 100 if total_tests > 0 else 0
    
    print(f"\n🎯 === FINAL SUMMARY ===")
    print(f"Total tests: {total_tests}")
    print(f"Successful: {successful_tests}")
    print(f"Success rate: {success_rate:.1f}%")
    
    print(f"\nBreakdown:")
    print(f"- Basic generation: {sum(1 for r in basic_results if r.get('success', False))}/{len(basic_results)}")
    print(f"- Conditioned generation: {sum(1 for r in conditioned_results if r.get('success', False))}/{len(conditioned_results)}")
    print(f"- Code completion: {sum(1 for r in code_results if r.get('success', False))}/{len(code_results)}")
    
    # Return all results for documentation
    return {
        'basic_results': basic_results,
        'conditioned_results': conditioned_results, 
        'code_results': code_results,
        'summary': {
            'total_tests': total_tests,
            'successful_tests': successful_tests,
            'success_rate': success_rate,
            'timestamp': datetime.now().isoformat()
        }
    }

if __name__ == "__main__":
    results = main()