import torch
from torch import nn
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models import resnet18
import segmentation_models_pytorch as smp
import torchvision
from pathlib import Path
import os

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"🖥️  Using device: {DEVICE}")

# Shared model registry - avoids Python import scoping issues with module globals
MODELS = {}

def build_frcnn():
    model = fasterrcnn_resnet50_fpn_v2(weights=None)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 33)
    return model

def build_classifier():
    model = resnet18(weights=None)
    model.fc = nn.Sequential(nn.Dropout(0.3), nn.Linear(model.fc.in_features, 32))
    return model

def build_unet():
    return smp.Unet(encoder_name="resnet34", encoder_weights=None, in_channels=3, classes=1)

def _strip_dp(state: dict) -> dict:
    """Strip 'module.' prefix added by nn.DataParallel on Kaggle multi-GPU."""
    if any(k.startswith("module.") for k in state.keys()):
        print("   ↳ Stripping DataParallel 'module.' prefix from checkpoint")
        return {k.replace("module.", "", 1): v for k, v in state.items()}
    return state

def load_convnext_model(ckpt_path: Path, device: torch.device):
    """Loads ConvNeXt-Tiny customized for 32 FDI classes (>90% Acc Setup)."""
    m = torchvision.models.convnext_tiny(weights=None)
    m.classifier[2] = nn.Linear(m.classifier[2].in_features, 32)
    
    if not os.path.exists(ckpt_path):
        print(f"⚠️  ConvNeXt Checkpoint not found: {ckpt_path}. Model will return random garbage until trained.")
        m.to(device)
        m.eval()
        return m
        
    state = torch.load(ckpt_path, map_location=device, weights_only=True)
    # Strip any "module." prefix if trained with DataParallel
    new_state = {}
    for k, v in state.items():
        name = k[7:] if k.startswith("module.") else k
        new_state[name] = v
        
    m.load_state_dict(new_state)
    m.to(device)
    m.eval()
    return m

def load_resnet_model(ckpt_path: Path, device: torch.device):
    """Loads ResNet18 customized for 32 FDI classes."""
    m = torchvision.models.resnet18(weights=None)
    m.fc = nn.Sequential(nn.Dropout(p=0.3), nn.Linear(m.fc.in_features, 32))
    state = torch.load(ckpt_path, map_location=device, weights_only=True)
    new_state = {}
    for k, v in state.items():
        name = k[7:] if k.startswith("module.") else k
        new_state[name] = v
    m.load_state_dict(new_state)
    m.to(device)
    m.eval()
    return m

def _load(path: str, model):
    if not os.path.exists(path):
        print(f"⚠️  Checkpoint not found: {path}")
        return model
    try:
        state = torch.load(path, map_location=DEVICE, weights_only=False)
        state = _strip_dp(state)
        model.load_state_dict(state, strict=False)
        print(f"✅ Loaded: {os.path.basename(path)}")
    except Exception as e:
        print(f"❌ Failed loading {os.path.basename(path)}: {e}")
    return model

def load_models(ckpt_dir: str):
    print(f"📦 Loading models from: {ckpt_dir}")
    MODELS["frcnn"] = _load(os.path.join(ckpt_dir, "frcnn_best.pt"),        build_frcnn()).to(DEVICE).eval()
    
    convnext_path = Path(os.path.join(ckpt_dir, "convnext_best.pt"))
    resnet_path = Path(os.path.join(ckpt_dir, "resnet18_cls_best.pt"))
    if convnext_path.exists():
        MODELS["cls"] = load_convnext_model(convnext_path, DEVICE)
        print("✅ ConvNeXt loaded.")
    elif resnet_path.exists():
        MODELS["cls"] = load_resnet_model(resnet_path, DEVICE)
        print("✅ ResNet18 loaded.")
    else:
        print("⚠️ No classification checkpoint found!")
        
    MODELS["unet"]  = _load(os.path.join(ckpt_dir, "unet_resnet34_best.pt"), build_unet()).to(DEVICE).eval()
    print("🚀 All PyTorch models ready.")