#!/usr/bin/env python3
"""
HarpoonNet 1.2 Model Loader
Utility to load HarpoonNet 1.2 with proper layer name mapping

Based on: https://huggingface.co/christiankhoury05/harpoon-1-2
"""

import torch
from harpoon_modular import create_harpoon_net_12

def load_harpoonnet_12(model_path='pytorch_model.pth', device='cpu'):
    """
    Load HarpoonNet 1.2 model with automatic layer name mapping
    
    Args:
        model_path: Path to the model checkpoint
        device: Device to load the model on ('cpu' or 'cuda')
        
    Returns:
        model: Loaded HarpoonNet 1.2 model ready for inference
    """
    print("🚀 Loading HarpoonNet 1.2 ConvNeXt model...")
    
    # Create model architecture
    model = create_harpoon_net_12(num_classes=1, num_anchors=3, pretrained=False)
    
    # Load checkpoint
    print(f"📁 Loading checkpoint from: {model_path}")
    checkpoint = torch.load(model_path, map_location='cpu')
    
    # Handle different checkpoint formats
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
        print("✅ Loaded from full checkpoint (model_state_dict)")
        
        # Print additional info if available
        if 'epoch' in checkpoint:
            print(f"   📊 Epoch: {checkpoint['epoch']}")
        if 'loss' in checkpoint:
            print(f"   📊 Training Loss: {checkpoint['loss']:.6f}")
        if 'val_loss' in checkpoint:
            print(f"   📊 Validation Loss: {checkpoint['val_loss']:.6f}")
    else:
        state_dict = checkpoint
        print("✅ Loaded weights-only checkpoint")
    
    # Map layer names from old architecture to new architecture
    print("🔧 Mapping layer names...")
    new_state_dict = {}
    
    for old_key, value in state_dict.items():
        # Map backbone.* -> features.*
        if old_key.startswith('backbone.'):
            new_key = old_key.replace('backbone.', 'features.')
            new_state_dict[new_key] = value
        # Keep detection_head layers as-is
        elif old_key.startswith('detection_head.'):
            new_state_dict[old_key] = value
        else:
            # Keep any other layers as-is
            new_state_dict[old_key] = value
    
    print(f"📊 Mapped {len(state_dict)} -> {len(new_state_dict)} layers")
    
    # Load the mapped state dict
    try:
        missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
        
        if missing_keys:
            print(f"⚠️  Missing keys: {len(missing_keys)}")
            for key in missing_keys[:3]:  # Show first 3
                print(f"     {key}")
            if len(missing_keys) > 3:
                print(f"     ... and {len(missing_keys) - 3} more")
        
        if unexpected_keys:
            print(f"⚠️  Unexpected keys: {len(unexpected_keys)}")
            for key in unexpected_keys[:3]:  # Show first 3
                print(f"     {key}")
            if len(unexpected_keys) > 3:
                print(f"     ... and {len(unexpected_keys) - 3} more")
        
        if not missing_keys and not unexpected_keys:
            print("✅ Perfect match! All layers loaded successfully")
        else:
            print("⚠️  Partial match - model may still work for inference")
        
    except Exception as e:
        print(f"❌ Failed to load state dict: {e}")
        return None
    
    # Move to device and set to eval mode
    model = model.to(device)
    model.eval()
    
    print(f"🎯 HarpoonNet 1.2 loaded successfully on {device}!")
    print("🔧 Model ready for inference")
    
    return model

def test_loaded_model(model, device='cpu'):
    """
    Test the loaded model with a dummy input
    
    Args:
        model: Loaded HarpoonNet 1.2 model
        device: Device the model is on
    """
    if model is None:
        print("❌ No model to test")
        return
    
    print("\n🧪 Testing loaded model...")
    
    # Create dummy input
    dummy_input = torch.randn(1, 3, 544, 544).to(device)
    
    # Test inference
    with torch.no_grad():
        output = model(dummy_input)
        detections = model.decode_predictions(output, confidence_threshold=0.5)
    
    print(f"✅ Inference test passed!")
    print(f"   Input shape: {dummy_input.shape}")
    print(f"   Output shape: {output.shape}")
    print(f"   Detections: {len(detections[0].get('boxes', []))}")

if __name__ == "__main__":
    # Test the loader
    print("🎯 HarpoonNet 1.2 Model Loader Test")
    print("=" * 50)
    
    # Load model
    model = load_harpoonnet_12('pytorch_model.pth', device='cpu')
    
    # Test it
    test_loaded_model(model, device='cpu')
    
    print("\n🎉 Model loader test complete!")
    print("📖 Usage:")
    print("   from load_harpoonnet_12 import load_harpoonnet_12")
    print("   model = load_harpoonnet_12('pytorch_model.pth')")