"""Construction and loading of ensemble models for evaluation.

Training saves each member's ``state_dict`` to ``model_N.pt``. Evaluation must
rebuild an identically-structured module before ``load_state_dict``. For Bayesian
members this means applying the DNN->BNN conversion (with the same prior params
used at training time) *before* loading weights, otherwise the parameter names
won't match.
"""

import pathlib as pl
from typing import Any, Dict

import torch
from torch import nn

from model.cnn import CNN3D
from model.dnn_mod import DEFAULT_BNN_PRIOR_PARAMETERS, dnn_to_bnn_mod

# Model-kind strings (mirror evaluation.schema.MODEL_KIND_*). Kept here too so
# model/ has no dependency on the evaluation package.
KIND_NORMAL = "normal"
KIND_BAYESIAN = "bayesian"


def build_model(config: Dict[str, Any], kind: str) -> CNN3D:
    """Build an (untrained) CNN3D of the requested kind on CPU.

    For ``kind == "bayesian"`` the conv/linear layers are converted to their
    Bayesian counterparts using :data:`DEFAULT_BNN_PRIOR_PARAMETERS`.
    """
    data_cfg = config["data"]
    train_cfg = config["training"]
    model = CNN3D(
        image_channels=data_cfg["image_channels"],
        clin_data_channels=data_cfg["clin_data_channels"],
        num_classes=data_cfg["num_classes"],
        droprate=train_cfg["droprate"],
    ).float()

    if kind == KIND_BAYESIAN:
        dnn_to_bnn_mod(model, DEFAULT_BNN_PRIOR_PARAMETERS)
    elif kind != KIND_NORMAL:
        raise ValueError(f"Unknown model kind {kind!r}; expected 'normal'/'bayesian'.")

    return model


def load_model(
    path: pl.Path,
    kind: str,
    config: Dict[str, Any],
    device: str,
) -> nn.Module:
    """Rebuild a model of ``kind`` and load ``path``'s weights onto ``device``.

    The returned model is in ``eval()`` mode. Bayesian models remain stochastic
    in eval mode (their reparameterization layers sample every forward pass),
    which is exactly what Monte-Carlo evaluation relies on.
    """
    model = build_model(config, kind)
    state_dict = torch.load(path, map_location="cpu")
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model