"""Swin-Unet wrapper.

Instantiates SwinTransformerSys from sota/Swin-Unet directly (bypassing the repo's
yacs config + .npz/.h5 Synapse pipeline). Handles grayscale by repeating 1->3.

Constraint: the windowed attention requires img_size = 224 (patches resolutions
56/28/14/7 are each divisible by window_size=7). Other sizes will assert; we keep
img_size=224 for this backbone.
"""
from __future__ import annotations

import os
import sys

import torch
import torch.nn as nn

_REPO = os.path.join(os.path.dirname(__file__), "..", "..", "sota", "Swin-Unet")
_REPO = os.path.abspath(_REPO)


def _ensure_path():
    if _REPO not in sys.path:
        sys.path.insert(0, _REPO)


class SwinUnetWrapper(nn.Module):
    def __init__(self, in_channels: int, num_classes: int, img_size: int = 224,
                 pretrained_ckpt: str = ""):
        super().__init__()
        _ensure_path()
        from networks.swin_transformer_unet_skip_expand_decoder_sys import SwinTransformerSys
        if img_size != 224:
            raise ValueError("Swin-Unet backbone requires img_size=224.")
        self.net = SwinTransformerSys(
            img_size=img_size, patch_size=4, in_chans=3, num_classes=num_classes,
            embed_dim=96, depths=[2, 2, 2, 2], num_heads=[3, 6, 12, 24],
            window_size=7, mlp_ratio=4.0, qkv_bias=True, drop_path_rate=0.1,
            ape=False, patch_norm=True, use_checkpoint=False,
        )
        if pretrained_ckpt and os.path.isfile(pretrained_ckpt):
            self._load_pretrained(pretrained_ckpt)

    def _load_pretrained(self, path):
        """Port of Swin-Unet's load_from: load the ImageNet Swin-T encoder AND
        mirror its `layers.X` weights into the decoder `layers_up.(3-X)` (the
        scheme the paper uses to initialize the symmetric decoder)."""
        import copy
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        if "model" not in ckpt:
            self.net.load_state_dict(ckpt, strict=False)
            return
        pretrained = ckpt["model"]
        model_dict = self.net.state_dict()
        full = copy.deepcopy(pretrained)
        for k, v in pretrained.items():
            if "layers." in k:
                n = 3 - int(k[7:8])
                full["layers_up." + str(n) + k[8:]] = v
        for k in list(full.keys()):
            if k in model_dict and full[k].shape != model_dict[k].shape:
                del full[k]
        msg = self.net.load_state_dict(full, strict=False)
        print(f"[swinunet] loaded pretrained {path}: "
              f"missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}")

    def forward(self, x):
        if x.size(1) == 1:
            x = x.repeat(1, 3, 1, 1)
        return self.net(x)


def build_swinunet(in_channels: int, num_classes: int, img_size: int = 224,
                   pretrained_ckpt: str = "", **_):
    return SwinUnetWrapper(in_channels, num_classes, img_size, pretrained_ckpt)