"""Детектор строк: компактный DB (differentiable binarization).

Пиксельная морфология упирается в пороги: один набор чисел разбирает газету и
рвёт книгу. Детектор учится на рамках строк, поэтому переносится на любую
вёрстку, лишь бы строка была строкой.

Почему именно DB, а не «просто сегментация текста»: в плотном газетном наборе
соседние строки и колонки сливаются в одно пятно. DB учит карту по *сжатым*
рамкам — между строками остаётся зазор — и отдельно учит порог, по которому
пятно снова разворачивается до полной строки.

Веса backbone не берутся предобученными: download.pytorch.org из пода
недоступен, а бесплатной разметки из born-digital PDF хватает для обучения
с нуля.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


def conv_bn(cin, cout, k=3, s=1):
    return nn.Sequential(
        nn.Conv2d(cin, cout, k, s, k // 2, bias=False),
        nn.BatchNorm2d(cout),
        nn.ReLU(inplace=True))


class Block(nn.Module):
    """Остаточный блок; понижение разрешения — первой свёрткой."""

    def __init__(self, cin, cout, stride=1):
        super().__init__()
        self.c1 = conv_bn(cin, cout, 3, stride)
        self.c2 = nn.Sequential(
            nn.Conv2d(cout, cout, 3, 1, 1, bias=False),
            nn.BatchNorm2d(cout))
        self.skip = None
        if stride != 1 or cin != cout:
            self.skip = nn.Sequential(
                nn.Conv2d(cin, cout, 1, stride, bias=False),
                nn.BatchNorm2d(cout))

    def forward(self, x):
        y = self.c2(self.c1(x))
        s = x if self.skip is None else self.skip(x)
        return F.relu(y + s, inplace=True)


class LineDet(nn.Module):
    """Вход: 1×H×W (0..1, чернила ближе к 1). Выход: карты в 1/2 разрешения.

    Шаг выхода 2, а не 4 как в оригинальном DB: в газете строка занимает
    10–14 px, при шаге 4 сжатая рамка вырождается в один пиксель и учить
    нечему.
    """

    def __init__(self, width=32):
        super().__init__()
        w = width
        self.stem = nn.Sequential(conv_bn(1, w, 3, 2), Block(w, w))          # 1/2
        self.s2 = nn.Sequential(Block(w, w * 2, 2), Block(w * 2, w * 2))     # 1/4
        self.s3 = nn.Sequential(Block(w * 2, w * 4, 2), Block(w * 4, w * 4))  # 1/8
        self.s4 = nn.Sequential(Block(w * 4, w * 8, 2), Block(w * 8, w * 8))  # 1/16

        self.l1 = nn.Conv2d(w, w * 2, 1)
        self.l2 = nn.Conv2d(w * 2, w * 2, 1)
        self.l3 = nn.Conv2d(w * 4, w * 2, 1)
        self.l4 = nn.Conv2d(w * 8, w * 2, 1)
        self.smooth = conv_bn(w * 2 * 4, w * 4)

        # две головы: вероятность текста и порог бинаризации
        self.prob = nn.Sequential(conv_bn(w * 4, w * 2), nn.Conv2d(w * 2, 1, 1))
        self.thr = nn.Sequential(conv_bn(w * 4, w * 2), nn.Conv2d(w * 2, 1, 1))
        self.k = 50.0                      # крутизна дифференцируемой бинаризации

    def forward(self, x):
        p1 = self.stem(x)
        p2 = self.s2(p1)
        p3 = self.s3(p2)
        p4 = self.s4(p3)
        c4 = self.l4(p4)
        c3 = self.l3(p3) + F.interpolate(c4, scale_factor=2, mode="nearest")
        c2 = self.l2(p2) + F.interpolate(c3, scale_factor=2, mode="nearest")
        c1 = self.l1(p1) + F.interpolate(c2, scale_factor=2, mode="nearest")
        f = torch.cat([
            c1,
            F.interpolate(c2, size=c1.shape[-2:], mode="nearest"),
            F.interpolate(c3, size=c1.shape[-2:], mode="nearest"),
            F.interpolate(c4, size=c1.shape[-2:], mode="nearest")], 1)
        f = self.smooth(f)
        prob = self.prob(f)
        thr = self.thr(f)
        return prob, thr

    def binary(self, prob, thr):
        """Дифференцируемая бинаризация: мягкая ступенька вместо сравнения."""
        return torch.sigmoid(self.k * (torch.sigmoid(prob) - torch.sigmoid(thr)))


def dice_loss(pred, target, mask, eps=1e-6):
    pred, target, mask = pred * mask, target * mask, mask
    inter = (pred * target).sum()
    return 1.0 - 2.0 * inter / (pred.sum() + target.sum() + eps)


def ohem_bce(logit, target, mask, ratio=3.0):
    """BCE только по всем положительным и втрое большему числу худших отрицательных.

    Текст занимает малую долю страницы; без отбора отрицательных обучение
    уезжает в «везде фон».
    """
    pos = (target > 0.5) & (mask > 0.5)
    neg = (target <= 0.5) & (mask > 0.5)
    n_pos = int(pos.sum())
    n_neg = min(int(neg.sum()), max(int(n_pos * ratio), 1000))
    loss = F.binary_cross_entropy_with_logits(logit, target, reduction="none")
    lp = loss[pos].sum() if n_pos else logit.sum() * 0.0
    ln = torch.topk(loss[neg], n_neg).values.sum() if n_neg else logit.sum() * 0.0
    return (lp + ln) / max(1, n_pos + n_neg)


def det_loss(model, prob_logit, thr_logit, gt_prob, gt_thr, mask_prob, mask_thr):
    l_prob = ohem_bce(prob_logit, gt_prob, mask_prob)
    l_thr = (torch.abs(torch.sigmoid(thr_logit) - gt_thr) * mask_thr).sum() \
        / mask_thr.sum().clamp(min=1)
    binary = model.binary(prob_logit, thr_logit)
    l_bin = dice_loss(binary, gt_prob, mask_prob)
    return l_prob + 10.0 * l_thr + l_bin, (l_prob.item(), l_thr.item(), l_bin.item())