#!/usr/bin/env python3
"""
CPU Debug Script for DFM_Qwen2VL Evaluation Pipeline

Purpose: Test the entire evaluation pipeline on CPU with minimal data
to verify no bugs before running expensive GPU evaluation.

Usage: python3 shiyi_scripts/debug_eval_cpu.py
"""

import os
import sys
import torch

# Set to CPU mode
os.environ['CUDA_VISIBLE_DEVICES'] = ''

# Add paths - need to add project root and chengyue_script
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
sys.path.insert(0, os.path.join(project_root, 'chengyue_script'))
sys.path.insert(0, os.path.join(project_root, 'src'))

# Also set PYTHONPATH for subprocess imports
os.environ['PYTHONPATH'] = f"{project_root}:{os.path.join(project_root, 'chengyue_script')}:{os.path.join(project_root, 'src')}:" + os.environ.get('PYTHONPATH', '')

print("=" * 60)
print("DFM_Qwen2VL CPU Debug Script")
print("=" * 60)
print()

# Step 1: Test model loading
print("Step 1: Testing model loading...")
print("-" * 60)

model_path = "src/models/Qwen2_5VL_DFM"
tokenizer_path = "Qwen/Qwen2.5-VL-3B-Instruct"

print(f"Model path: {model_path}")
print(f"Tokenizer path: {tokenizer_path}")
print()

try:
    from transformers import AutoTokenizer, AutoModelForCausalLM
    
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_path,
        trust_remote_code=True
    )
    print(f"✅ Tokenizer loaded: vocab_size={tokenizer.vocab_size}")
    print()
    
    print("Loading model (this may take a while)...")
    print("Note: Loading Qwen2.5-VL-3B base weights on CPU...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        torch_dtype=torch.float32,  # Use float32 on CPU
        device_map='cpu'
    )
    print(f"✅ Model loaded successfully")
    print(f"   Model type: {type(model).__name__}")
    print(f"   Number of parameters: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")
    print()
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Step 2: Test tokenization
print("Step 2: Testing tokenization...")
print("-" * 60)

test_texts = [
    "What is 2 + 2?",
    "Write a Python function to add two numbers.",
]

try:
    for i, text in enumerate(test_texts):
        tokens = tokenizer(text, return_tensors="pt")
        print(f"Text {i+1}: {text}")
        print(f"  Input IDs shape: {tokens['input_ids'].shape}")
        print(f"  Num tokens: {tokens['input_ids'].shape[1]}")
    print()
    print("✅ Tokenization working correctly")
    print()
    
except Exception as e:
    print(f"❌ Error in tokenization: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Step 3: Test forward pass
print("Step 3: Testing forward pass...")
print("-" * 60)

try:
    model.eval()
    
    test_input = "Hello, how are you?"
    print(f"Test input: {test_input}")
    
    inputs = tokenizer(test_input, return_tensors="pt")
    input_ids = inputs['input_ids']
    
    print(f"Input shape: {input_ids.shape}")
    
    with torch.no_grad():
        outputs = model(input_ids)
        logits = outputs.logits
    
    print(f"Output logits shape: {logits.shape}")
    print(f"✅ Forward pass successful")
    print()
    
except Exception as e:
    print(f"❌ Error in forward pass: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Step 4: Test generation (simple greedy decoding)
print("Step 4: Testing simple generation...")
print("-" * 60)

try:
    prompt = "What is 2 + 2? Answer:"
    print(f"Prompt: {prompt}")
    
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs['input_ids']
    
    print(f"Generating 10 tokens with greedy decoding...")
    
    with torch.no_grad():
        for step in range(10):
            outputs = model(input_ids)
            logits = outputs.logits
            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
            input_ids = torch.cat([input_ids, next_token], dim=1)
            
            if step < 5:  # Print first few steps
                decoded = tokenizer.decode(next_token[0])
                print(f"  Step {step+1}: token_id={next_token[0].item()}, decoded='{decoded}'")
    
    generated_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
    print()
    print(f"Generated text: {generated_text}")
    print()
    print("✅ Simple generation working")
    print()
    
except Exception as e:
    print(f"❌ Error in generation: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Step 5: Test DFM-specific components
print("Step 5: Testing DFM-specific components...")
print("-" * 60)

try:
    # Check if model has DFM attributes
    if hasattr(model, 'model'):
        dfm_model = model.model
        print(f"✅ Found DFM model: {type(dfm_model).__name__}")
        
        if hasattr(dfm_model, 'bd_size'):
            print(f"   bd_size: {dfm_model.bd_size}")
        if hasattr(dfm_model, 'a'):
            print(f"   a parameter: {dfm_model.a}")
        if hasattr(dfm_model, 'c'):
            print(f"   c parameter: {dfm_model.c}")
    else:
        print("⚠️  Model doesn't have 'model' attribute (might be expected)")
    
    # Check config
    if hasattr(model, 'config'):
        config = model.config
        print(f"✅ Model config:")
        print(f"   model_type: {config.model_type}")
        if hasattr(config, 'bd_size'):
            print(f"   bd_size: {config.bd_size}")
        if hasattr(config, 'base_model_name_or_path'):
            print(f"   base_model: {config.base_model_name_or_path}")
    
    print()
    print("✅ DFM components accessible")
    print()
    
except Exception as e:
    print(f"⚠️  Warning in DFM component check: {e}")
    print("   (This might be OK if model structure is different)")
    print()

# Step 6: Test batch processing
print("Step 6: Testing batch processing...")
print("-" * 60)

try:
    batch_texts = [
        "What is 1 + 1?",
        "What is 2 + 2?",
        "What is 3 + 3?",
    ]
    
    print(f"Processing batch of {len(batch_texts)} texts...")
    
    # Tokenize with padding
    inputs = tokenizer(
        batch_texts,
        return_tensors="pt",
        padding=True,
        truncation=True
    )
    
    print(f"Batch input shape: {inputs['input_ids'].shape}")
    
    with torch.no_grad():
        outputs = model(inputs['input_ids'])
        logits = outputs.logits
    
    print(f"Batch output shape: {logits.shape}")
    print("✅ Batch processing working")
    print()
    
except Exception as e:
    print(f"❌ Error in batch processing: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Step 7: Test with mask tokens (if DFM model)
print("Step 7: Testing with mask tokens (DFM-specific)...")
print("-" * 60)

try:
    # Get mask token ID
    mask_id = 151665  # Default mask ID for Qwen
    
    # Test if mask token is in vocabulary
    if mask_id < tokenizer.vocab_size:
        print(f"Mask token ID: {mask_id}")
        
        # Create input with some mask tokens
        test_ids = torch.tensor([[1, 2, mask_id, mask_id, 5, 6]])
        print(f"Input with masks: {test_ids}")
        
        with torch.no_grad():
            outputs = model(test_ids)
            logits = outputs.logits
        
        print(f"Output shape: {logits.shape}")
        print("✅ Model can handle mask tokens")
    else:
        print(f"⚠️  Mask ID {mask_id} >= vocab_size {tokenizer.vocab_size}")
    
    print()
    
except Exception as e:
    print(f"⚠️  Warning in mask token test: {e}")
    print()

# Summary
print("=" * 60)
print("SUMMARY")
print("=" * 60)
print()
print("✅ All basic tests passed!")
print()
print("Pipeline components verified:")
print("  ✅ Model loading")
print("  ✅ Tokenization")
print("  ✅ Forward pass")
print("  ✅ Simple generation")
print("  ✅ DFM components")
print("  ✅ Batch processing")
print("  ✅ Mask token handling")
print()
print("Next steps:")
print("  1. Model can be loaded successfully")
print("  2. Basic inference works on CPU")
print("  3. Ready for GPU evaluation once training completes")
print()
print("To run full evaluation on GPU with checkpoint:")
print("  1. Wait for training checkpoint")
print("  2. Update model_path in run_batch_eval_dvlm_3b_slurm.sh")
print("  3. sbatch shiyi_scripts/run_batch_eval_dvlm_3b_slurm.sh")
print()
print("=" * 60)

