"""Satr — Xattot's text-line recognizer (PyTorch definition, for fine-tuning).

Inference does not need this file: `satr/satr.onnx` runs with onnxruntime only.
Architecture: CNN (height 48 -> 1, width / 4) + 2 Transformer blocks (d=192, 4 heads) + linear head, trained with CTC.
3.26M parameters. Classes are copied unchanged from the training notebook, so the weights load with strict=True.

    import torch, satr_model
    model, chars = satr_model.load("satr/satr.safetensors", "satr/vocab.json")
    x = torch.from_numpy(line_gray_uint8[None, None])        # (1, 1, 48, W), white background, height 48
    print(satr_model.decode(model(x).log_softmax(-1), chars))

Input: grayscale line image resized to height 48 (keep aspect ratio). CTC blank = index 0, class i = chars[i - 1].
License: CC BY-NC-SA 4.0 (see LICENSE). Fine-tuned weights may be shared under the same license with attribution.
"""
import json

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


class Blok(nn.Module):
    def __init__(self, ci, co, s):
        super().__init__()
        self.c1 = nn.Conv2d(ci, co, 3, s, 1, bias=False); self.b1 = nn.BatchNorm2d(co)
        self.c2 = nn.Conv2d(co, co, 3, 1, 1, bias=False); self.b2 = nn.BatchNorm2d(co)
        self.sk = nn.Sequential(nn.Conv2d(ci, co, 1, s, bias=False), nn.BatchNorm2d(co)) if (ci != co or s != 1) else nn.Identity()

    def forward(self, x):
        return F.gelu(self.b2(self.c2(F.gelu(self.b1(self.c1(x))))) + self.sk(x))


class TBlok(nn.Module):
    # own Transformer block: nn.MultiheadAttention fixes the width in TorchScript ONNX export; shapes here stay dynamic
    def __init__(self, d, h=4, ff=3, p=0.1):
        super().__init__()
        self.h, self.p = h, p
        self.n1, self.qkv, self.o = nn.LayerNorm(d), nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.n2, self.ff = nn.LayerNorm(d), nn.Sequential(nn.Linear(d, d * ff), nn.GELU(), nn.Linear(d * ff, d))
        self.dr = nn.Dropout(p)

    def forward(self, x):                                            # (N,T,d)
        q, k, v = self.qkv(self.n1(x)).unflatten(-1, (3, self.h, -1)).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v, dropout_p=self.p if self.training else 0.0)
        x = x + self.dr(self.o(a.transpose(1, 2).flatten(2)))
        return x + self.dr(self.ff(self.n2(x)))


class Satr(nn.Module):
    def __init__(self, n_cls, d=192):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1, 1, bias=False), nn.BatchNorm2d(32), nn.GELU(),
            Blok(32, 48, (2, 2)), Blok(48, 48, 1),                    # 24 x W/2
            Blok(48, 96, (2, 2)), Blok(96, 96, 1),                    # 12 x W/4
            Blok(96, 144, (2, 1)), Blok(144, 144, 1),                 # 6 x W/4
            Blok(144, d, (2, 1)), Blok(d, d, 1))                      # 3 x W/4
        self.yig = nn.Conv2d(d, d, (3, 1), bias=False)                # 1 x W/4
        self.pos = nn.Conv1d(d, d, 9, padding=4, groups=d)
        self.tr = nn.Sequential(TBlok(d), TBlok(d))
        self.out = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, n_cls))

    def forward(self, x):                                            # x: (N,1,48,W) uint8, or float already normalized to -1..1
        x = (x.float() / 255.0 - 0.5) / 0.5 if x.dtype == torch.uint8 else x
        f = self.yig(self.cnn(x)).squeeze(2)                         # (N,d,T)
        f = f + self.pos(f)
        return self.out(self.tr(f.transpose(1, 2)))                  # (N,T,C) logits


def load(weights="satr/satr.safetensors", vocab="satr/vocab.json", device="cpu"):
    from safetensors.torch import load_file
    chars = json.load(open(vocab, encoding="utf-8"))["chars"]
    model = Satr(len(chars) + 1)
    model.load_state_dict(load_file(weights, device="cpu"), strict=True)
    return model.to(device).eval(), chars


def decode(logp, chars):
    """Greedy CTC decoding of (N,T,C) scores."""
    out = []
    for s in logp.argmax(-1).cpu().numpy():
        p, t = 0, []
        for c in s:
            if c != p and c != 0:
                t.append(chars[c - 1])
            p = c
        out.append("".join(t))
    return out