"""Original reduced U-Net with valid convolutions and cropped skip features.

Two resolution reductions, widths 4/8/16, input 96x96 and output 56x56.
The original 2015 network is deeper and wider. No trained segmentation weights.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3)
        self.relu2 = nn.ReLU()

    def forward(self, x):
        return self.relu2(self.conv2(self.relu1(self.conv1(x))))


class CropAndJoin(nn.Module):
    def forward(self, decoder, skip):
        top = (skip.shape[-2] - decoder.shape[-2]) // 2
        left = (skip.shape[-1] - decoder.shape[-1]) // 2
        cropped = skip[
            ..., top : top + decoder.shape[-2], left : left + decoder.shape[-1]
        ]
        return torch.cat([decoder, cropped], dim=1)


# tensorviz: input=1,1,96,96
class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder1 = DoubleConv(1, 4)
        self.pool1 = nn.MaxPool2d(2)
        self.encoder2 = DoubleConv(4, 8)
        self.pool2 = nn.MaxPool2d(2)
        self.bottleneck = DoubleConv(8, 16)
        self.up2 = nn.ConvTranspose2d(16, 8, 2, stride=2)
        self.join2 = CropAndJoin()
        self.decoder2 = DoubleConv(16, 8)
        self.up1 = nn.ConvTranspose2d(8, 4, 2, stride=2)
        self.join1 = CropAndJoin()
        self.decoder1 = DoubleConv(8, 4)
        self.mask_head = nn.Conv2d(4, 2, 1)

    def forward(self, x):
        skip1 = self.encoder1(x)
        skip2 = self.encoder2(self.pool1(skip1))
        bridge = self.bottleneck(self.pool2(skip2))
        coarse = self.decoder2(self.join2(self.up2(bridge), skip2))
        fine = self.decoder1(self.join1(self.up1(coarse), skip1))
        return self.mask_head(fine)


def verify():
    torch.manual_seed(10)
    model = TinyUNet().double().eval()
    x = torch.randn(1, 1, 96, 96, dtype=torch.float64)
    with torch.no_grad():
        assert model(x).shape == (1, 2, 56, 56)
        skip1 = model.encoder1(x)
        skip2 = model.encoder2(model.pool1(skip1))
        bridge = model.bottleneck(model.pool2(skip2))
        up = model.up2(bridge)
        joined = model.join2(up, skip2)
        assert joined.shape == (1, 16, 34, 34)
        torch.testing.assert_close(joined[:, :8], up, rtol=0, atol=0)
        torch.testing.assert_close(
            joined[:, 8:], skip2[..., 4:38, 4:38], rtol=0, atol=0
        )
    decoder = torch.randn(1, 2, 4, 4, dtype=torch.float64, requires_grad=True)
    skip = torch.randn(1, 3, 8, 8, dtype=torch.float64, requires_grad=True)
    joined = CropAndJoin()(decoder, skip)
    joined[:, 2:].sum().backward()
    expected = torch.zeros_like(skip)
    expected[..., 2:6, 2:6] = 1
    torch.testing.assert_close(skip.grad, expected, rtol=0, atol=0)
    torch.testing.assert_close(decoder.grad, torch.zeros_like(decoder), rtol=0, atol=0)
    return [
        "Valid convolutions map 96×96 inputs to two-channel 56×56 logits",
        "Upsampling to 34×34 joins eight decoder and eight centrally cropped skip channels",
        "The crop preserves exact skip values, rather than interpolating them",
        "A skip-only loss sends gradients through the cropped region and not through the decoder branch",
    ]


def thumbnail(tensor):
    return F.adaptive_avg_pool2d(tensor[0, :1], (8, 8))[0].tolist()


def experiment():
    torch.manual_seed(10)
    model = TinyUNet().double().eval()
    x = torch.zeros(1, 1, 96, 96, dtype=torch.float64)
    x[..., 24:72, 32:64] = 1
    cases = []
    with torch.no_grad():
        s1 = model.encoder1(x)
        s2 = model.encoder2(model.pool1(s1))
        bridge = model.bottleneck(model.pool2(s2))
        u2 = model.up2(bridge)
        j2 = model.join2(u2, s2)
        d2 = model.decoder2(j2)
        u1 = model.up1(d2)
        j1 = model.join1(u1, s1)
        for identity, label, up, skip, joined in [
            ("coarse", "Coarse skip: 42 → 34 pixels", u2, s2, j2),
            ("fine", "Fine skip: 92 → 60 pixels", u1, s1, j1),
        ]:
            channels = up.shape[1]
            cropped = joined[:, channels:]
            cases.append(
                {
                    "id": identity,
                    "label": label,
                    "target": "join2" if identity == "coarse" else "join1",
                    "note": "Each square below is an average-pooled thumbnail of one recorded feature channel. The skip is center-cropped to match the upsampled decoder; concatenation keeps both channel sets. These randomly initialized features are not a learned segmentation or an explanation of pixel importance.",
                    "matrices": [
                        {
                            "label": "Decoder channel 0 · 8×8 pooled thumbnail",
                            "values": thumbnail(up),
                        },
                        {
                            "label": "Skip before crop · 8×8 pooled thumbnail",
                            "values": thumbnail(skip),
                        },
                        {
                            "label": "Skip after crop · 8×8 pooled thumbnail",
                            "values": thumbnail(cropped),
                        },
                    ],
                    "vectors": [
                        {
                            "label": "Decoder / skip / joined channel counts",
                            "values": [channels, skip.shape[1], joined.shape[1]],
                        }
                    ],
                    "metrics": [
                        {
                            "label": "Skip spatial width before crop",
                            "value": skip.shape[-1],
                        },
                        {"label": "Shared width after crop", "value": up.shape[-1]},
                        {
                            "label": "Pixels cropped from each side",
                            "value": (skip.shape[-1] - up.shape[-1]) // 2,
                        },
                    ],
                }
            )
    return {
        "kind": "matrices",
        "title": "Recover detail through cropped skip features.",
        "description": "Recorded activations from a synthetic rectangle image. U-Net joins high-resolution encoder features to the expanding decoder along the channel axis. Thumbnails show spatial structure; exact dimensions and channel counts are separate.",
        "controlLabel": "Skip connection",
        "cases": cases,
    }