"""
Loading the released weights.

Every file under weights/ is a plain dict of tensors + primitives, so it loads
under `torch.load(..., weights_only=True)` -- no custom classes, no pickled
training objects, no need to have this repo importable to read it.

Layout of a checkpoint file::

    {
      "state_dict":   OrderedDict[str, Tensor],   # no "module." prefix
      "arch":         "c3d" | "mednext" | "swinunetr",
      "in_channels":  8 (pretrained) or 5 (finetuned / fromscratch),
      "out_channels": 1,
      "anatomy":      "han" | "pancreas" | "source-domain",
      "regime":       "pretrained" | "finetuned" | "fromscratch",
      "act_sig":      True,        # apply sigmoid to the raw output
      "scale_out":    7.5 | 5.5,   # then multiply by this
      "dose_div_factor": 10,       # then multiply by this -> dose in Gy
      "source": "...", "note": "...",
    }

So the raw network output becomes physical dose as::

    dose_Gy = sigmoid(out) * scale_out * dose_div_factor
"""

import torch

from build_model import build_model


def load_checkpoint(path, map_location="cpu"):
    """Read a released checkpoint file. Returns the full dict."""
    return torch.load(path, map_location=map_location, weights_only=True)


def load_model(path, device="cpu", strict=True):
    """
    Build the right architecture for a checkpoint and load its weights.

    Returns ``(model_in_eval_mode, checkpoint_dict)``. The checkpoint dict
    carries the post-processing constants (``act_sig``, ``scale_out``,
    ``dose_div_factor``) you need to turn the output into Gy.
    """
    ckpt = load_checkpoint(path, map_location="cpu")

    model = build_model(
        ckpt["arch"],
        in_channels=ckpt["in_channels"],
        out_channels=ckpt.get("out_channels", 1),
    )
    model.load_state_dict(ckpt["state_dict"], strict=strict)
    model.to(device).eval()
    return model, ckpt


def to_dose(output, ckpt):
    """
    Convert a raw network output to dose in Gy using the checkpoint's own
    constants. Handles C3D's ``[output_A, output_B]`` and MedNeXt's optional
    deep-supervision list by taking the refined / main head.
    """
    if isinstance(output, (list, tuple)):
        # C3D: [coarse A, refined B] -> take B. MedNeXt DS: [main, aux...] -> take main.
        output = output[-1] if ckpt["arch"] == "c3d" else output[0]

    if ckpt.get("act_sig", True):
        output = torch.sigmoid(output)

    return output * ckpt["scale_out"] * ckpt["dose_div_factor"]