"""3D ResNet over the 8-channel atlas volume (4 MR types + 4 presence channels)."""
from __future__ import annotations

import math

import torch.nn as nn


def _norm(c: int) -> nn.Module:
    # small per-GPU batches: GroupNorm with the largest group count <= 32 that divides c
    return nn.GroupNorm(math.gcd(c, 32), c)


class BasicBlock(nn.Module):
    def __init__(self, cin: int, cout: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv3d(cin, cout, 3, stride, 1, bias=False)
        self.n1 = _norm(cout)
        self.conv2 = nn.Conv3d(cout, cout, 3, 1, 1, bias=False)
        self.n2 = _norm(cout)
        self.act = nn.SiLU(inplace=True)
        self.down = None
        if stride != 1 or cin != cout:
            self.down = nn.Sequential(nn.Conv3d(cin, cout, 1, stride, bias=False), _norm(cout))

    def forward(self, x):
        idt = x if self.down is None else self.down(x)
        x = self.act(self.n1(self.conv1(x)))
        x = self.n2(self.conv2(x))
        return self.act(x + idt)


class ResNet3D(nn.Module):
    """(B, 8, D, H, W) -> (B, num_classes) logits. Total stride 16 (stride-2 stem + three stride-2 stages)."""

    def __init__(self, in_ch: int = 8, num_classes: int = 74, width: int = 48,
                 blocks: tuple[int, ...] = (2, 2, 2, 2), dropout: float = 0.2):
        super().__init__()
        w = width
        self.stem = nn.Sequential(
            nn.Conv3d(in_ch, w, 3, 2, 1, bias=False), _norm(w), nn.SiLU(inplace=True),
            nn.Conv3d(w, w, 3, 1, 1, bias=False), _norm(w), nn.SiLU(inplace=True),
        )
        stages, cin = [], w
        for i, (cout, n) in enumerate(zip([w, w * 2, w * 4, w * 8], blocks)):
            layers = [BasicBlock(cin, cout, stride=2 if i > 0 else 1)]
            layers += [BasicBlock(cout, cout) for _ in range(n - 1)]
            stages.append(nn.Sequential(*layers))
            cin = cout
        self.stages = nn.Sequential(*stages)
        self.head = nn.Sequential(nn.AdaptiveAvgPool3d(1), nn.Flatten(), nn.Dropout(dropout),
                                  nn.Linear(cin, num_classes))
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")

    def forward(self, x):
        return self.head(self.stages(self.stem(x)))


ARCHS = {
    "r18": dict(width=48, blocks=(2, 2, 2, 2)),
    "r18w": dict(width=64, blocks=(2, 2, 2, 2)),
    "r26w": dict(width=64, blocks=(2, 3, 4, 2)),
}


def build_model(arch: str = "r18", **kw) -> ResNet3D:
    return ResNet3D(**ARCHS[arch], **kw)