model_loader.py

"""
model_loader.py
Downloads weights from HF_WEIGHTS_REPO once, loads all models into memory.

Key guarantees
──────────────
1. PYEOF removed (was a SyntaxError that silently broke everything).
2. Sentinel file (.load_ok) written after first successful download.
   On container restart the sentinel + file-size checks bypass the
   download entirely so only the torch.load / model.eval() work remains.
3. calibration.pkl is unpickled via three independent strategies so it
   never raises even when saved from a different Python __main__ context.
4. Thread-safe: only one goroutine ever calls load_all_models().
"""

import logging, os, pickle, sys, threading, types
from pathlib import Path
from typing import Optional, Tuple

import torch
from huggingface_hub import hf_hub_download

import model_classes
from model_classes import CalibrationBundle, DRModel, FundusValidator, NUM_CLASSES

logger = logging.getLogger("model_loader")

# ─────────────────────────────────────────────────────────────────────────────
# Patch sys.modules["__main__"] at import time — must happen before any pickle
# ─────────────────────────────────────────────────────────────────────────────
def _patch_main_now():
    main = sys.modules.get("__main__")
    if main is not None:
        for _name in ["CalibrationBundle", "GeM", "CoralHead", "DRHead",
                      "DRModel", "FundusValidator"]:
            if not hasattr(main, _name):
                setattr(main, _name, getattr(model_classes, _name, None))
    shim = types.ModuleType("__main__")
    for _name in dir(model_classes):
        setattr(shim, _name, getattr(model_classes, _name))
    cur = sys.modules.get("__main__")
    if cur is not None and not hasattr(cur, "CalibrationBundle"):
        sys.modules["__main__"] = shim


_patch_main_now()   # runs at import time

# ─────────────────────────────────────────────────────────────────────────────
# Custom Unpickler — triple-checks every find_class call
# ─────────────────────────────────────────────────────────────────────────────
_CLASS_MAP = {
    "CalibrationBundle": model_classes.CalibrationBundle,
    "GeM":               model_classes.GeM,
    "CoralHead":         model_classes.CoralHead,
    "DRHead":            model_classes.DRHead,
    "DRModel":           model_classes.DRModel,
    "FundusValidator":   model_classes.FundusValidator,
}

class _Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if name in _CLASS_MAP:
            return _CLASS_MAP[name]
        if module in ("__main__", "__mp_main__"):
            cls = getattr(model_classes, name, None)
            if cls is not None:
                return cls
        return super().find_class(module, name)


# ─────────────────────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────────────────────
HF_WEIGHTS_REPO: str    = os.environ.get("HF_WEIGHTS_REPO", "Gokul-G1/dr-grading-weights")
HF_TOKEN: Optional[str] = os.environ.get("HF_TOKEN", None)

# /data is HF Persistent Storage (survives container restarts).
# Fall back to /tmp only when no persistent storage is attached.
_STORAGE_ROOT: Path = (
    Path("/data/dr_models") if Path("/data").exists() else Path("/tmp/dr_models")
)
_STORAGE_ROOT.mkdir(parents=True, exist_ok=True)

# Minimum file sizes (bytes) — reject partial/corrupt downloads
_MIN_SIZES = {
    "final_complete_model.pt": 50 * 1024 * 1024,   # >50 MB
    "fundus_validator.pt":      5 * 1024 * 1024,   #  >5 MB
    "calibration.pkl":          1 * 1024,           #  >1 KB
}
_SENTINEL = _STORAGE_ROOT / ".load_ok"

# ─────────────────────────────────────────────────────────────────────────────
# Singletons
# ─────────────────────────────────────────────────────────────────────────────
_dr_model:    Optional[DRModel]           = None
_fundus_val:  Optional[FundusValidator]   = None
_calibration: Optional[CalibrationBundle] = None
_device:      str                         = "cpu"
_lock                                     = threading.Lock()
_ready:       bool                        = False


def _pick_device() -> str:
    if torch.cuda.is_available():
        logger.info(f"[loader] GPU: {torch.cuda.get_device_name(0)}")
        return "cuda"
    logger.info("[loader] No GPU — using CPU")
    return "cpu"


def _file_ok(name: str) -> bool:
    """True if file exists on disk and passes minimum-size check."""
    p = _STORAGE_ROOT / name
    if not p.exists():
        return False
    size = p.stat().st_size
    min_size = _MIN_SIZES.get(name, 1)
    if size < min_size:
        logger.warning(f"[cache] {name} too small ({size} B < {min_size} B) — will re-download")
        p.unlink(missing_ok=True)
        return False
    return True


def _all_files_ok() -> bool:
    return all(_file_ok(n) for n in _MIN_SIZES)


def _download(filename: str) -> Path:
    """Download only if not already cached on disk."""
    if _file_ok(filename):
        local = _STORAGE_ROOT / filename
        logger.info(f"[cache] {filename}  ({local.stat().st_size/1e6:.1f} MB) — using cached copy")
        return local
    logger.info(f"[download] {filename} from {HF_WEIGHTS_REPO} …")
    # FIX: local_dir_use_symlinks may be removed in newer huggingface_hub — graceful fallback
    try:
        path = hf_hub_download(
            repo_id=HF_WEIGHTS_REPO, filename=filename,
            local_dir=str(_STORAGE_ROOT), token=HF_TOKEN,
            local_dir_use_symlinks=False,
        )
    except TypeError:
        path = hf_hub_download(
            repo_id=HF_WEIGHTS_REPO, filename=filename,
            local_dir=str(_STORAGE_ROOT), token=HF_TOKEN,
        )
    local = Path(path)
    if not local.parent.samefile(_STORAGE_ROOT):
        # hf_hub_download may have placed it in a subdir — copy flat
        dest = _STORAGE_ROOT / filename
        dest.write_bytes(local.read_bytes())
        local = dest
    logger.info(f"[download] done → {local}  ({local.stat().st_size/1e6:.1f} MB)")
    return local


def _torch_load(path: Path, map_location="cpu"):
    try:
        return torch.load(str(path), map_location=map_location, weights_only=False)
    except TypeError:
        return torch.load(str(path), map_location=map_location)


def _load_calibration(cal_path: Path) -> Optional[CalibrationBundle]:
    """
    Try every known strategy to load calibration.pkl.
    Returns None if all fail — inference continues with raw softmax.
    """
    _patch_main_now()   # belt and braces

    # Strategy 1: custom _Unpickler (handles __main__ pickle artifacts)
    try:
        with open(cal_path, "rb") as f:
            obj = _Unpickler(f).load()
        logger.info(f"[loader] calibration loaded via _Unpickler  type={type(obj).__name__}")
        return obj
    except Exception as e1:
        logger.warning(f"[loader] _Unpickler failed: {e1}")

    # Strategy 2: torch.load (notebook may have used torch.save on a plain dict)
    try:
        obj = _torch_load(cal_path, map_location="cpu")
        logger.info(f"[loader] calibration loaded via torch.load  type={type(obj).__name__}")
        return obj
    except Exception as e2:
        logger.warning(f"[loader] torch.load failed: {e2}")

    # Strategy 3: plain pickle with force-patched __main__
    try:
        cur = sys.modules.get("__main__")
        if cur is not None:
            cur.CalibrationBundle = model_classes.CalibrationBundle
        with open(cal_path, "rb") as f:
            obj = pickle.load(f)
        logger.info(f"[loader] calibration loaded via plain pickle  type={type(obj).__name__}")
        return obj
    except Exception as e3:
        logger.warning(f"[loader] plain pickle failed: {e3}")

    logger.error("[loader] All calibration strategies failed — inference runs without calibration")
    return None


# ─────────────────────────────────────────────────────────────────────────────
# Main loader — idempotent, thread-safe
# ─────────────────────────────────────────────────────────────────────────────
def load_all_models() -> None:
    global _dr_model, _fundus_val, _calibration, _device, _ready
    with _lock:
        if _ready:
            return

        device  = _pick_device()
        _device = device
        logger.info(f"[loader] device={device}  repo={HF_WEIGHTS_REPO}  storage={_STORAGE_ROOT}")

        # Check sentinel: if it exists all files are already on disk, skip network
        if _SENTINEL.exists() and _all_files_ok():
            logger.info("[loader] Sentinel present + all files OK — skipping download step")
        else:
            logger.info("[loader] Downloading model artefacts …")
            _SENTINEL.unlink(missing_ok=True)   # invalidate stale sentinel

        # ── Grading model ─────────────────────────────────────────────────────
        dr_path = _download("final_complete_model.pt")
        ckpt    = _torch_load(dr_path, map_location="cpu")

        if isinstance(ckpt, dict):
            backbone = ckpt.get("backbone", "tf_efficientnetv2_m")
            state    = (ckpt.get("model_state")
                        or ckpt.get("state_dict")
                        or ckpt.get("ema_state_dict")
                        or ckpt)
        else:
            backbone = "tf_efficientnetv2_m"
            state    = ckpt

        model = DRModel(backbone=backbone, num_classes=NUM_CLASSES, pretrained=False)
        miss, unexpected = model.load_state_dict(state, strict=False)
        if miss:
            # Filter out known non-critical keys
            real_miss = [k for k in miss if not k.startswith("_")]
            logger.warning(f"[loader] missing keys ({len(real_miss)}): {real_miss[:8]}")
        if unexpected:
            logger.warning(f"[loader] unexpected keys ({len(unexpected)}): {unexpected[:5]}")
        model.eval().to(device)
        _dr_model = model
        n = sum(p.numel() for p in model.parameters()) / 1e6
        logger.info(f"[loader] DRModel ready  {n:.1f}M params  backbone={backbone}")

        # ── Fundus validator ──────────────────────────────────────────────────
        try:
            fv_path = _download("fundus_validator.pt")
            fv_ckpt = _torch_load(fv_path, map_location="cpu")

            # Robustly extract state dict regardless of how it was saved
            if isinstance(fv_ckpt, dict):
                # Try every common wrapper key in order
                fv_state = (
                    fv_ckpt.get("model_state")
                    or fv_ckpt.get("state_dict")
                    or fv_ckpt.get("model")
                    or fv_ckpt.get("net")
                    or None
                )
                # If none of the wrapper keys matched, the dict IS the state dict
                # — but only if it has parameter-shaped values (tensors)
                if fv_state is None:
                    first_val = next(iter(fv_ckpt.values()), None)
                    if isinstance(first_val, torch.Tensor):
                        fv_state = fv_ckpt
                    else:
                        raise ValueError(
                            f"Cannot find state dict in fundus_validator.pt "
                            f"(top-level keys: {list(fv_ckpt.keys())[:8]})"
                        )
            else:
                # checkpoint IS the state dict (OrderedDict of tensors)
                fv_state = fv_ckpt

            fv = FundusValidator(pretrained=False)
            missing, unexpected = fv.load_state_dict(fv_state, strict=False)
            if missing:
                logger.warning(f"[loader] FundusValidator missing keys ({len(missing)}): "
                               f"{missing[:5]}")
            if unexpected:
                logger.warning(f"[loader] FundusValidator unexpected keys ({len(unexpected)}): "
                               f"{unexpected[:5]}")
            fv.eval().to(device)
            _fundus_val = fv
            logger.info("[loader] FundusValidator ready")
        except Exception as fv_err:
            logger.error(f"[loader] FundusValidator load FAILED: {fv_err} "
                         "— heuristic-only validation will be used", exc_info=True)
            _fundus_val = None  # inference.py handles None gracefully

        # ── Calibration bundle ────────────────────────────────────────────────
        cal_path     = _download("calibration.pkl")
        _calibration = _load_calibration(cal_path)   # never raises

        # Write sentinel so next restart skips download
        _SENTINEL.write_text("ok")
        _ready = True
        cal_status = type(_calibration).__name__ if _calibration else "SKIPPED (None)"
        logger.info(f"[loader] ✓ All models ready  calibration={cal_status}")


def get_models() -> Tuple[DRModel, FundusValidator, Optional[CalibrationBundle], str]:
    if not _ready:
        load_all_models()
    return _dr_model, _fundus_val, _calibration, _device