"""Custom 3D residual U-Net (~10.7M params), valid-padding throughout.

Valid (unpadded) convolutions are used instead of same-padding because this setup is intended to run
under blockwise inference (daisy/gunpowder tiling adjacent crops across a large volume). Same-padding
convs synthesize zero-context at every block's border, which produces seam artifacts at block
boundaries; valid convs give every output voxel genuine surrounding context, matching the convention
already used by fly_organelles.model.StandardUnet and fish_trainer.model.Unet3d.

Verified (base_channels=40, factor=2, blocks_per_level=1): 10,702,493 params,
input (148,148,148) -> output (60,60,60).
"""

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


def center_crop_to(x, target):
    diffs = [x.shape[i] - target.shape[i] for i in (2, 3, 4)]
    starts = [d // 2 for d in diffs]
    return x[
        :,
        :,
        starts[0] : starts[0] + target.shape[2],
        starts[1] : starts[1] + target.shape[3],
        starts[2] : starts[2] + target.shape[4],
    ]


class ResBlock(nn.Module):
    """Two valid 3x3x3 convs + InstanceNorm + ReLU, with a crop-and-add identity skip
    (1x1x1 conv on the skip path when channel count changes; identity is center-cropped
    to match the main path's spatially-shrunk output before the residual add)."""

    def __init__(self, c_in, c_out):
        super().__init__()
        self.conv1 = nn.Conv3d(c_in, c_out, 3, padding=0)
        self.norm1 = nn.InstanceNorm3d(c_out, affine=True)
        self.conv2 = nn.Conv3d(c_out, c_out, 3, padding=0)
        self.norm2 = nn.InstanceNorm3d(c_out, affine=True)
        self.act = nn.ReLU(inplace=True)
        self.skip = nn.Conv3d(c_in, c_out, 1) if c_in != c_out else nn.Identity()

    def forward(self, x):
        idn = self.skip(x)
        h = self.act(self.norm1(self.conv1(x)))
        h = self.norm2(self.conv2(h))
        idn = center_crop_to(idn, h)
        return self.act(h + idn)


class Down(nn.Module):
    """Exact 2x downsample (kernel==stride, divides evenly, no cropping needed)."""

    def __init__(self, c_in, c_out):
        super().__init__()
        self.conv = nn.Conv3d(c_in, c_out, 2, stride=2)

    def forward(self, x):
        return self.conv(x)


class Up(nn.Module):
    """Non-learnable nearest-neighbor upsample (exact 2x) + 1x1x1 channel-reduce conv.
    Deliberately non-learnable to match this codebase's own constant_upsample=True choice in
    StandardUnet (avoids checkerboard/seam artifacts from learned transposed convs)."""

    def __init__(self, c_in, c_out):
        super().__init__()
        self.reduce = nn.Conv3d(c_in, c_out, 1)

    def forward(self, x):
        x = F.interpolate(x, scale_factor=2, mode="nearest")
        return self.reduce(x)


class ResUNet3DValid(nn.Module):
    def __init__(self, out_channels, base_channels=40, factor=2, blocks_per_level=1, in_channels=1):
        super().__init__()
        C = [base_channels * (factor**i) for i in range(4)]  # 3 down/up levels + bottleneck
        self.C = C

        self.stem = ResBlock(in_channels, C[0])

        self.downs = nn.ModuleList()
        self.enc_blocks = nn.ModuleList()
        for i in range(3):
            self.downs.append(Down(C[i], C[i + 1]))
            self.enc_blocks.append(nn.ModuleList([ResBlock(C[i + 1], C[i + 1]) for _ in range(blocks_per_level)]))

        self.ups = nn.ModuleList()
        self.dec_blocks = nn.ModuleList()
        for i in reversed(range(3)):
            self.ups.append(Up(C[i + 1], C[i]))
            blocks = [ResBlock(C[i] * 2, C[i])] + [ResBlock(C[i], C[i]) for _ in range(blocks_per_level - 1)]
            self.dec_blocks.append(nn.ModuleList(blocks))

        self.final_conv = nn.Conv3d(C[0], out_channels, 1)

    def forward(self, input):
        skips = []
        x = self.stem(input)
        skips.append(x)
        for down, enc in zip(self.downs, self.enc_blocks):
            x = down(x)
            for blk in enc:
                x = blk(x)
            skips.append(x)
        skips.pop()  # bottleneck itself is not a skip target

        for up, dec, skip in zip(self.ups, self.dec_blocks, reversed(skips)):
            x = up(x)
            skip_c = center_crop_to(skip, x)
            x = torch.cat([x, skip_c], dim=1)
            for blk in dec:
                x = blk(x)

        return self.final_conv(x)