"""Standalone PyTorch 3D model definitions and loader for Vesuvius Surface Segmentation."""

import torch
import torch.nn as nn
import torch.nn.functional as F


class ResidualBlock3D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.norm1 = nn.InstanceNorm3d(out_channels, affine=True)
        self.act1 = nn.LeakyReLU(negative_slope=0.01, inplace=True)
        self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.norm2 = nn.InstanceNorm3d(out_channels, affine=True)
        self.act2 = nn.LeakyReLU(negative_slope=0.01, inplace=True)
        if in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=False),
                nn.InstanceNorm3d(out_channels, affine=True),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        res = self.shortcut(x)
        x = self.act1(self.norm1(self.conv1(x)))
        x = self.norm2(self.conv2(x))
        return self.act2(x + res)


class DecoderBlock3D(nn.Module):
    def __init__(self, in_channels: int, skip_channels: int, out_channels: int):
        super().__init__()
        self.upsample = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2, bias=False)
        self.norm_up = nn.InstanceNorm3d(in_channels // 2, affine=True)
        self.act_up = nn.LeakyReLU(negative_slope=0.01, inplace=True)
        self.conv = ResidualBlock3D((in_channels // 2) + skip_channels, out_channels)

    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
        x = self.act_up(self.norm_up(self.upsample(x)))
        if x.shape[2:] != skip.shape[2:]:
            x = F.interpolate(x, size=skip.shape[2:], mode="trilinear", align_corners=False)
        return self.conv(torch.cat([x, skip], dim=1))


class ResEncUNet3D(nn.Module):
    """5-stage 3D U-Net with Residual Encoder blocks (~15.91M parameters)."""

    def __init__(self, in_channels: int = 1, out_channels: int = 1, channels=(32, 64, 128, 256, 320)):
        super().__init__()
        c0, c1, c2, c3, c4 = channels
        self.enc0 = ResidualBlock3D(in_channels, c0)
        self.down0 = nn.MaxPool3d(2, 2)
        self.enc1 = ResidualBlock3D(c0, c1)
        self.down1 = nn.MaxPool3d(2, 2)
        self.enc2 = ResidualBlock3D(c1, c2)
        self.down2 = nn.MaxPool3d(2, 2)
        self.enc3 = ResidualBlock3D(c2, c3)
        self.down3 = nn.MaxPool3d(2, 2)
        self.bottleneck = ResidualBlock3D(c3, c4)
        self.dec3 = DecoderBlock3D(c4, c3, c3)
        self.dec2 = DecoderBlock3D(c3, c2, c2)
        self.dec1 = DecoderBlock3D(c2, c1, c1)
        self.dec0 = DecoderBlock3D(c1, c0, c0)
        self.out_conv = nn.Conv3d(c0, out_channels, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        e0 = self.enc0(x)
        e1 = self.enc1(self.down0(e0))
        e2 = self.enc2(self.down1(e1))
        e3 = self.enc3(self.down2(e2))
        b = self.bottleneck(self.down3(e3))
        u3 = self.dec3(b, e3)
        u2 = self.dec2(u3, e2)
        u1 = self.dec1(u2, e1)
        u0 = self.dec0(u1, e0)
        return self.out_conv(u0)


class _SegConvWrap(nn.Module):
    def __init__(self, in_c: int, out_c: int, k: int = 3, s: int = 1, p: int = 1, bias: bool = False):
        super().__init__()
        self.conv = nn.Conv3d(in_c, out_c, kernel_size=k, stride=s, padding=p, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv(x)


class _SegResBlock(nn.Module):
    def __init__(self, c: int):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, c)
        self.norm2 = nn.GroupNorm(8, c)
        self.act = nn.ReLU(inplace=True)
        self.conv1 = _SegConvWrap(c, c, 3, 1, 1, False)
        self.conv2 = _SegConvWrap(c, c, 3, 1, 1, False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.conv2(self.act(self.norm2(self.conv1(self.act(self.norm1(x))))))


class _SegUpSampleWrap(nn.Module):
    def __init__(self):
        super().__init__()
        self.upsample_non_trainable = nn.Upsample(scale_factor=(2.0, 2.0, 2.0), mode="trilinear", align_corners=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.upsample_non_trainable(x)


class PureSegResNet3D(nn.Module):
    """Pure PyTorch SegResNet 3D (~18.80M parameters, bit-for-bit MONAI SegResNet state_dict compatible)."""

    def __init__(self, in_channels: int = 1, out_channels: int = 1, init_filters: int = 32, blocks_down=(1, 2, 2, 4), blocks_up=(1, 1, 1)):
        super().__init__()
        self.act_mod = nn.ReLU(inplace=True)
        self.convInit = _SegConvWrap(in_channels, init_filters, 3, 1, 1, False)
        self.down_layers = nn.ModuleList()
        c = init_filters
        for i, n_blk in enumerate(blocks_down):
            layers = [nn.Identity() if i == 0 else _SegConvWrap(c // 2, c, 3, 2, 1, False)]
            layers.extend(_SegResBlock(c) for _ in range(n_blk))
            self.down_layers.append(nn.Sequential(*layers))
            if i < len(blocks_down) - 1:
                c *= 2
        self.up_layers = nn.ModuleList()
        self.up_samples = nn.ModuleList()
        for _, n_blk in enumerate(blocks_up):
            half_c = c // 2
            self.up_samples.append(nn.Sequential(_SegConvWrap(c, half_c, 1, 1, 0, False), _SegUpSampleWrap()))
            self.up_layers.append(nn.Sequential(*[_SegResBlock(half_c) for _ in range(n_blk)]))
            c = half_c
        self.conv_final = nn.Sequential(nn.GroupNorm(8, c), nn.ReLU(inplace=True), _SegConvWrap(c, out_channels, 1, 1, 0, True))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.convInit(x)
        skips = []
        for down in self.down_layers:
            x = down(x)
            skips.append(x)
        skips.pop()
        for up_s, up_l in zip(self.up_samples, self.up_layers):
            x = up_s(x) + skips.pop()
            x = up_l(x)
        return self.conv_final(x)


def load_model(architecture: str, checkpoint_path: str, device: str = "cpu") -> nn.Module:
    """Instantiate architecture ('PureSegResNet3D' or 'ResEncUNet3D') and load checkpoint weights."""
    if architecture == "PureSegResNet3D":
        model = PureSegResNet3D()
    elif architecture == "ResEncUNet3D":
        model = ResEncUNet3D()
    else:
        raise ValueError(f"Unsupported architecture: {architecture}")

    state = torch.load(checkpoint_path, map_location=device, weights_only=True)
    if isinstance(state, dict) and ("model_state_dict" in state or "state_dict" in state):
        state = state.get("model_state_dict", state.get("state_dict"))
    model.load_state_dict(state)
    model.to(device).eval()
    return model