"""Text-segment detector: MobileNetV2-style backbone + FPN + segmentation head.

Predicts two probability maps at input resolution:
  0. line kernels: each text segment's central region (half its height, with its ends trimmed by a
     quarter height). Post-processing expands each kernel back to the full box, so adjacent lines stay
     separated even when their line spacing is tight.
  1. block kernels: each layout block (paragraph, list item, heading, table, caption...) as the rectangle
     around its lines shrunk by BLOCK_SHRINK line heights on every side, so neighbouring blocks stay apart.
     Lines are grouped into paragraphs by the block region their centre falls in.
"""

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

from .common import ConvBNAct

KERNEL_HEIGHT_RATIO = 0.5   # kernel height = ratio * box height
KERNEL_END_TRIM = 0.25      # kernel width = box width - 2 * trim * box height
BLOCK_SHRINK = 0.3          # block kernel = block rectangle shrunk by this many line heights per side


class InvertedResidual(nn.Module):
    def __init__(self, cin, cout, stride=1, expand=4):
        super().__init__()
        mid = cin * expand
        self.use_res = stride == 1 and cin == cout
        self.body = nn.Sequential(
            ConvBNAct(cin, mid, k=1),
            ConvBNAct(mid, mid, k=3, s=stride, groups=mid),
            ConvBNAct(mid, cout, k=1, act=False),
        )

    def forward(self, x):
        y = self.body(x)
        return x + y if self.use_res else y


def _stage(cin, cout, n, stride, expand):
    blocks = [InvertedResidual(cin, cout, stride, expand)]
    blocks += [InvertedResidual(cout, cout, 1, expand) for _ in range(n - 1)]
    return nn.Sequential(*blocks)


class Detector(nn.Module):
    def __init__(self, widths=(24, 48, 104, 192), depths=(2, 3, 4, 3), expand=4, fpn=64):
        super().__init__()
        self.stem = ConvBNAct(3, 16, s=2)                                  # /2
        self.s1 = _stage(16, widths[0], depths[0], 2, expand)              # /4
        self.s2 = _stage(widths[0], widths[1], depths[1], 2, expand)       # /8
        self.s3 = _stage(widths[1], widths[2], depths[2], 2, expand)       # /16
        self.s4 = _stage(widths[2], widths[3], depths[3], 2, expand)       # /32
        self.lateral = nn.ModuleList(nn.Conv2d(w, fpn, 1) for w in widths)
        self.smooth = nn.ModuleList(ConvBNAct(fpn, fpn // 4) for _ in widths)
        self.head = nn.Sequential(
            ConvBNAct(fpn, 64),
            nn.ConvTranspose2d(64, 32, 2, 2), nn.BatchNorm2d(32), nn.ReLU(inplace=True),  # /2
            nn.ConvTranspose2d(32, 2, 2, 2),                                               # /1: lines, blocks
        )

    def forward(self, x):
        """x: (B, 3, H, W) in [0, 1], H and W multiples of 32. Returns logits (B, 2, H, W): line, block kernels."""
        x = self.stem(x - 0.5)
        c2 = self.s1(x)
        c3 = self.s2(c2)
        c4 = self.s3(c3)
        c5 = self.s4(c4)
        p5 = self.lateral[3](c5)
        p4 = self.lateral[2](c4) + F.interpolate(p5, size=c4.shape[-2:], mode="nearest")
        p3 = self.lateral[1](c3) + F.interpolate(p4, size=c3.shape[-2:], mode="nearest")
        p2 = self.lateral[0](c2) + F.interpolate(p3, size=c2.shape[-2:], mode="nearest")
        size = p2.shape[-2:]
        fused = torch.cat([
            self.smooth[0](p2),
            F.interpolate(self.smooth[1](p3), size=size, mode="nearest"),
            F.interpolate(self.smooth[2](p4), size=size, mode="nearest"),
            F.interpolate(self.smooth[3](p5), size=size, mode="nearest"),
        ], dim=1)
        return self.head(fused)