"""
model.py — U-Net for radio coverage map prediction.

Input : [B, 2, H, W]  (buildings + transmitter location)
Output: [B, 1, H, W]  (predicted coverage map, values in [0, 1])

Classic encoder-decoder with skip connections (Ronneberger et al., 2015),
adapted for image-to-image regression instead of segmentation.
"""

import torch
import torch.nn as nn


class DoubleConv(nn.Module):
    """(conv 3x3 -> BatchNorm -> ReLU) applied twice."""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

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


class Down(nn.Module):
    """Halve resolution (max-pool) then DoubleConv."""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.pool = nn.MaxPool2d(2)
        self.conv = DoubleConv(in_ch, out_ch)

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


class Up(nn.Module):
    """Upsample, concatenate the skip connection, then DoubleConv."""
    def __init__(self, in_ch, skip_ch, out_ch):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_ch // 2 + skip_ch, out_ch)

    def forward(self, x, skip):
        x = self.up(x)
        x = torch.cat([skip, x], dim=1)   # skip connection
        return self.conv(x)


class UNet(nn.Module):
    def __init__(self, in_channels=2, out_channels=1, base=64):
        super().__init__()
        self.inc   = DoubleConv(in_channels, base)   # 256
        self.down1 = Down(base,     base * 2)        # 128
        self.down2 = Down(base * 2, base * 4)        # 64
        self.down3 = Down(base * 4, base * 8)        # 32
        self.down4 = Down(base * 8, base * 8)        # 16  (bottleneck)

        self.up1 = Up(base * 8, base * 8, base * 4)
        self.up2 = Up(base * 4, base * 4, base * 2)
        self.up3 = Up(base * 2, base * 2, base)
        self.up4 = Up(base,     base,     base)

        self.outc = nn.Conv2d(base, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        return torch.sigmoid(self.outc(x))


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = UNet(in_channels=2, out_channels=1, base=64).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Device: {device}")
    print(f"Parameters: {n_params:,}")

    x = torch.randn(4, 2, 256, 256, device=device)
    y = model(x)
    print("Input :", tuple(x.shape))
    print("Output:", tuple(y.shape), "| range:", round(float(y.min()), 3), "-", round(float(y.max()), 3))
    if device == "cuda":
        print(f"Peak VRAM (forward, batch 4): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")