"""
One place that turns an architecture name + input-channel count into a model.

The three architectures are built exactly as they were during training, so the
released state_dicts load with `strict=True`. If you change any argument here
(feature_size, model_id, list_ch_*, ...) the checkpoints will no longer fit.

MedNeXt and SwinUNETR are imported lazily so that, e.g., running the C3D model
does not require the MedNeXt package to be installed.
"""

import torch.nn as nn

# --- architecture arguments, frozen to match the released checkpoints -------
C3D_KWARGS = dict(
    list_ch_A=[-1, 16, 32, 64, 128, 256],
    list_ch_B=[-1, 32, 64, 128, 256, 512],
)

MEDNEXT_KWARGS = dict(
    model_id="B",
    kernel_size=3,
    deep_supervision=False,
)

SWINUNETR_KWARGS = dict(
    spatial_dims=3,
    feature_size=96,            # "L" size
    dropout_path_rate=0.1,
    use_v2=True,
    use_checkpoint=True,
)

# These models were trained with MONAI < 1.3, whose SwinUNETR took an `img_size`
# argument. It was only used to sanity-check the window size and never affected
# the parameters, and MONAI removed it in 1.5. We pass it only if the installed
# version still accepts it, so one checkpoint works across MONAI versions.
SWINUNETR_IMG_SIZE = (128, 192, 192)   # Z, Y, X

ARCHITECTURES = ("c3d", "mednext", "swinunetr")


def build_model(arch: str, in_channels: int = 5, out_channels: int = 1) -> nn.Module:
    """
    Build an untrained model.

    arch         : one of ARCHITECTURES
    in_channels  : 5 for every fine-tuned / from-scratch checkpoint,
                   8 for the source-domain *pretrained* checkpoints
    out_channels : 1 (dose)

    Note the C3D forward returns a list ``[output_A, output_B]``; ``output_B``
    is the refined prediction and the one you want. MedNeXt and SwinUNETR
    return a single tensor.
    """
    arch = arch.lower()

    if arch == "c3d":
        from model_c3d import Model
        return Model(in_ch=in_channels, out_ch=out_channels, **C3D_KWARGS)

    if arch == "mednext":
        # pip install git+https://github.com/MIC-DKFZ/MedNeXt.git
        from nnunet_mednext import create_mednext_v1
        return create_mednext_v1(
            num_input_channels=in_channels,
            num_classes=out_channels,
            **MEDNEXT_KWARGS,
        )

    if arch == "swinunetr":
        import inspect
        from monai.networks.nets import SwinUNETR

        kwargs = dict(SWINUNETR_KWARGS)
        if "img_size" in inspect.signature(SwinUNETR.__init__).parameters:
            kwargs["img_size"] = SWINUNETR_IMG_SIZE
        return SwinUNETR(
            in_channels=in_channels,
            out_channels=out_channels,
            **kwargs,
        )

    raise ValueError(f"unknown arch {arch!r}; expected one of {ARCHITECTURES}")