"""Text-line recognizer: CNN feature extractor -> Transformer encoder -> CTC head.

Input is a grayscale line crop of height 32 and variable width. The CNN collapses height to 1 and
downsamples width by 4, so each output step covers 4 input columns. Position comes from a depthwise
convolution (relative, translation-equivariant) rather than an absolute table, so lines longer than
any seen in training are handled the same way as short ones.
"""

import torch
from torch import nn

from ..charset import NUM_CLASSES
from .common import ConvBNAct, TransformerBlock, padding_bias

LINE_HEIGHT = 32
WIDTH_STRIDE = 4
X_STRETCH = 1.25  # widen crops so condensed fonts still give CTC >= 1 step per character (+ blanks)


def _separable(cin, cout):
    """Depthwise 3x3 then pointwise 1x1: about a ninth of a full 3x3 convolution's arithmetic."""
    return nn.Sequential(ConvBNAct(cin, cin, groups=cin), ConvBNAct(cin, cout, k=1))


class Recognizer(nn.Module):
    """cnn="full" (v1-v4) or "separable": the same shape of front-end with depthwise-separable convolutions,
    measured 2.5x (4 threads) to 3x (1 thread) faster per page on CPU, where the full CNN is ~70% of the time."""

    def __init__(self, num_classes=NUM_CLASSES, d=192, layers=5, heads=6, ff=672, dropout=0.1, cnn="full"):
        super().__init__()
        conv = ConvBNAct if cnn == "full" else _separable
        self.cnn = nn.Sequential(
            ConvBNAct(1, 32), nn.MaxPool2d(2, 2),                       # 16 x W/2
            conv(32, 64), nn.MaxPool2d(2, 2),                           # 8 x W/4
            conv(64, 128), conv(128, 128), nn.MaxPool2d((2, 1), (2, 1)),   # 4 x W/4
            conv(128, 192), conv(192, 192), nn.MaxPool2d((2, 1), (2, 1)),  # 2 x W/4
            ConvBNAct(192, 256, k=(2, 3), p=(0, 1)),                    # 1 x W/4
        )
        self.proj = nn.Linear(256, d)
        self.pos = nn.Conv1d(d, d, 7, padding=3, groups=d)
        self.blocks = nn.ModuleList(TransformerBlock(d, heads, ff, dropout) for _ in range(layers))
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, num_classes)

    def forward(self, images, step_mask=None):
        """images: (B, 1, 32, W) in [0, 1], ink dark. step_mask: optional (B, W/4) bool, True = padding.

        Returns logits (B, T, C) with T = W / 4.
        """
        f = self.cnn(images - 0.5).squeeze(2).transpose(1, 2)  # (B, T, 256)
        x = self.proj(f)
        if step_mask is not None:  # keep padding from leaking into real steps through the conv
            x = x.masked_fill(step_mask[..., None], 0.0)
        x = x + self.pos(x.transpose(1, 2)).transpose(1, 2)
        bias = None if step_mask is None else padding_bias(step_mask, x.dtype)
        for blk in self.blocks:
            x = blk(x, bias)
        return self.head(self.ln(x))