"""
Builders for the four architectures we compare.

Loaded via importlib so the parent dirs of the original model files are not
inserted into sys.path (otherwise their per-model `dataset.py` would shadow
this experiment's `dataset.py`).

Models:
    segnet         — pv_panel_models/cnn_model/cnn_segmenter.py        (SegNet)
    unet           — pv_panel_models/unet_model/unet_model.py          (U-Net)
    segformer_b0   — pv_panel_models/vit_model/segformer_model.py      (SegFormer mit-b0)
    segformer_b5   — pv_panel_models/segformer_b5_model/segformer_model.py  (SegFormer mit-b5)

NOTE: SegNet's `forward()` already applies sigmoid; UNet/SegFormer return raw
logits. The trainer uses `output_is_prob=True` for SegNet's metrics step.

SegNet's loss is reproduced inline (BCELoss + Dice on probabilities) since
pv_panel_models/cnn_model/train.py uses a sibling-relative import that doesn't
survive being loaded by importlib without a sys.path tweak.
"""
import importlib.util
from pathlib import Path

import torch
import torch.nn as nn

REPO_ROOT = Path(__file__).resolve().parents[2]
PV_DIR = REPO_ROOT / "pv_panel_models"


def _load(module_name: str, file_path: Path):
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"could not load {file_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_segnet_mod = _load("_pv_segnet_model", PV_DIR / "cnn_model" / "cnn_segmenter.py")
_unet_mod = _load("_pv_unet_model", PV_DIR / "unet_model" / "unet_model.py")
_segformer_b0_mod = _load("_pv_segformer_b0_model", PV_DIR / "vit_model" / "segformer_model.py")
_segformer_b5_mod = _load("_pv_segformer_b5_model",
                          PV_DIR / "segformer_b5_model" / "segformer_model.py")


# SegNet expects probabilities (its forward applies sigmoid).
# Mirrors pv_panel_models/cnn_model/train.py:CombinedLoss exactly.
class _SegNetDiceLoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        pred = pred.view(-1)
        target = target.view(-1)
        intersection = (pred * target).sum()
        dice = (2.0 * intersection + self.smooth) / (pred.sum() + target.sum() + self.smooth)
        return 1 - dice


class _SegNetCombinedLoss(nn.Module):
    def __init__(self, bce_weight=0.5):
        super().__init__()
        self.bce = nn.BCELoss()  # SegNet output is already in [0,1]
        self.dice = _SegNetDiceLoss()
        self.bce_weight = bce_weight

    def forward(self, pred, target):
        return self.bce_weight * self.bce(pred, target) + (1 - self.bce_weight) * self.dice(pred, target)


def build_segnet():
    model = _segnet_mod.SegNet(in_channels=3, out_channels=1)
    loss = _SegNetCombinedLoss(bce_weight=0.5)
    return model, loss, True  # output_is_prob (sigmoid in forward)


def build_unet():
    model = _unet_mod.UNet(in_channels=3, out_channels=1)
    loss = _unet_mod.CombinedLoss(bce_weight=0.5)
    return model, loss, False


def build_segformer_b0():
    model = _segformer_b0_mod.SegformerModel(pretrained_name="nvidia/mit-b0", num_classes=1)
    loss = _segformer_b0_mod.CombinedLoss(bce_weight=0.5)
    return model, loss, False


def build_segformer_b5():
    model = _segformer_b5_mod.SegformerModel(pretrained_name="nvidia/mit-b5", num_classes=1)
    loss = _segformer_b5_mod.CombinedLoss(bce_weight=0.5)
    return model, loss, False


MODEL_REGISTRY = {
    "segnet":       build_segnet,
    "unet":         build_unet,
    "segformer_b0": build_segformer_b0,
    "segformer_b5": build_segformer_b5,
}

PRETTY_NAME = {
    "segnet":       "SegNet (CNN)",
    "unet":         "U-Net",
    "segformer_b0": "SegFormer-B0",
    "segformer_b5": "SegFormer-B5",
}