import os
import sys
import torch

# Add the parent directory of this script to the python path
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from model.config import VaghLMConfig
from model.transformer import VaghLM

def load_and_inspect_model():
    # 1. Load default configuration
    config = VaghLMConfig()
    print("--- Model Configuration ---")
    print(f"Model Name:         {config.model_name}")
    print(f"Vocabulary Size:    {config.vocab_size}")
    print(f"Number of Layers:   {config.n_layers}")
    print(f"Attention Heads:    {config.n_heads}")
    print(f"Embedding Dim:      {config.n_embd}")
    print(f"Feed-Forward Dim:   {config.ffn_dim}")
    print(f"Context Length:     {config.context_length}")
    print(f"Weight Tying:       {config.weight_tying}")
    print("-" * 30)

    # 2. Instantiate the model
    print("\nInstantiating VaghLM model...")
    model = VaghLM(config)
    
    # 3. Print parameter count details
    param_count = model.count_parameters()
    print(f"Total Model Parameters: {param_count:,} ({param_count / 1e6:.2f}M)")

    # 4. Load a checkpoint (if available)
    checkpoint_dir = os.path.join(ROOT, "checkpoints")
    checkpoint_names = ["finetuned_v2_best.pt", "best_checkpoint.pt"]
    checkpoint_loaded = False

    for name in checkpoint_names:
        ckpt_path = os.path.join(checkpoint_dir, name)
        if os.path.exists(ckpt_path):
            print(f"\nLoading weights from: {ckpt_path}")
            # Map location to CPU to avoid CUDA dependency for inspection
            checkpoint = torch.load(ckpt_path, map_location="cpu")
            
            # Load the state dict
            model.load_state_dict(checkpoint["model_state"])
            checkpoint_loaded = True
            
            print("[OK] Weights loaded successfully!")
            print(f"Checkpoint Metadata:")
            print(f"  - Step: {checkpoint.get('step', 'N/A')}")
            print(f"  - Validation Loss: {checkpoint.get('loss', 'N/A')}")
            break

    if not checkpoint_loaded:
        print("\n[INFO] No checkpoint files found. The model is initialized with random weights.")

    # 5. Run a dummy forward pass to verify functionality
    print("\nRunning test forward pass...")
    model.eval()
    
    # Generate some dummy token IDs (batch_size=2, sequence_length=10)
    # Ensure token IDs are within vocab_size bounds [0, 31999]
    dummy_input = torch.randint(low=0, high=config.vocab_size, size=(2, 10))
    print(f"Input Shape: {dummy_input.shape}")
    
    with torch.no_grad():
        logits, loss = model(dummy_input)
        
    print(f"Logits Shape: {logits.shape} (Expected: [batch_size, sequence_length, vocab_size])")
    print(f"Output Vocab Size: {logits.size(-1)}")
    print("[SUCCESS] Forward pass test successful!")

if __name__ == "__main__":
    load_and_inspect_model()