"""ResNet34 backbone + 3-scale FPN-style fusion.

Reproduces the paper's genuinely distinctive idea (multi-scale down/up-sample fusion
to help extract thin, elongated lane structure) with a concretely-defined FPN instead
of the paper's ambiguous "32-group ResNet32/ResNeXt50" description.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torchvision


class ConvBNMish(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, padding: int = 1):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(out_ch)
        self.act = nn.Mish(inplace=True)

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


class Backbone(nn.Module):
    """Fuses stride 8/16/32 features (paper's 3-scale idea), then downsamples the
    fused map to stride 32 before returning it -- full O(N^2) self-attention over a
    stride-8 map (e.g. 45x80=3600 tokens) is prohibitively expensive/OOMs on CPU;
    the paper's own tiny FLOPs count (0.425 GMACs) implies the same heavy
    downsampling before attention, just left unstated. Stride 32 on a 360x640 input
    gives an 11x20=220-token sequence, ~260x cheaper attention.
    """

    def __init__(self, name: str = "resnet34", pretrained: bool = True, out_channels: int = 128):
        super().__init__()
        weights = torchvision.models.ResNet34_Weights.IMAGENET1K_V1 if pretrained else None
        net = torchvision.models.resnet34(weights=weights)

        self.stem = nn.Sequential(net.conv1, net.bn1, net.relu, net.maxpool)
        self.layer1 = net.layer1  # stride 4
        self.layer2 = net.layer2  # stride 8,  C=128
        self.layer3 = net.layer3  # stride 16, C=256
        self.layer4 = net.layer4  # stride 32, C=512

        # Shared extractor delta: project each scale to a common channel dim.
        self.reduce_c3 = ConvBNMish(128, out_channels, kernel_size=1, padding=0)
        self.reduce_c4 = ConvBNMish(256, out_channels, kernel_size=1, padding=0)
        self.reduce_c5 = ConvBNMish(512, out_channels, kernel_size=1, padding=0)

        self.fuse_c4 = ConvBNMish(out_channels, out_channels)
        self.fuse_c3 = ConvBNMish(out_channels, out_channels)

        # stride 8 -> stride 32 for the transformer input (see class docstring).
        self.downsample = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
        )

        self.out_channels = out_channels

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.layer1(x)
        c3 = self.layer2(x)      # stride 8
        c4 = self.layer3(c3)     # stride 16
        c5 = self.layer4(c4)     # stride 32

        p5 = self.reduce_c5(c5)
        p4 = self.reduce_c4(c4) + nn.functional.interpolate(p5, size=c4.shape[-2:], mode="nearest")
        p4 = self.fuse_c4(p4)
        p3 = self.reduce_c3(c3) + nn.functional.interpolate(p4, size=c3.shape[-2:], mode="nearest")
        p3 = self.fuse_c3(p3)

        return self.downsample(p3)  # (B, out_channels, H/32, W/32)