"""Inference entry point for the exported CADC model repository."""
from pathlib import Path

import numpy as np
import torch

from cadc.model import build_model


def load_ensemble(root, device="cpu"):
    models = []
    names = []
    for run, arch in (("baseline", "resnet3d"), ("multiview", "multiview2d")):
        for fold in range(5):
            path = Path(root) / "runs" / run / f"fold{fold}" / "final.pt"
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)
            model = build_model(arch, pretrained=False)
            model.load_state_dict(checkpoint["model"])
            models.append(model.eval().to(device))
            names.append(f"{arch}/fold{fold}")
    return models, names


@torch.inference_mode()
def predict_patch(models, patch, device="cpu"):
    """Score one 64^3 HU patch, indexed (z, y, x), at 1 mm voxel spacing.

    The central 48^3 crop and HU scaling match cadc.data.NoduleDataset.
    The returned score estimates radiologist suspicion, not cancer probability.
    """
    patch = np.asarray(patch)
    if patch.shape != (64, 64, 64):
        raise ValueError("Expected one (64, 64, 64) nodule patch at 1 mm spacing.")
    if patch.dtype.kind not in "iuf" or not np.isfinite(patch).all():
        raise ValueError("Patch must contain finite numeric Hounsfield units.")
    if not models:
        raise ValueError("No models loaded.")
    crop = np.clip(patch[8:56, 8:56, 8:56].astype(np.float32), -1000, 400)
    x = torch.from_numpy(np.ascontiguousarray((crop + 1000) / 1400))[None, None].to(device)
    votes = [float(torch.sigmoid(model(x)).item()) for model in models]
    return float(np.mean(votes)), votes