"""Field and relationship model over a page's text segments.

Each segment is embedded from its characters, geometry and shape flags. A small Transformer with a
learned 2D spatial attention bias mixes context across the page, then three heads predict:
  * a field/structure label per segment,
  * pairwise relations (same_row, key_value),
  * the document type (receipt / invoice / general).
"""

import torch
from torch import nn

from ..charset import EXTRACTOR_CLASSES
from ..features import NUM_FEATS
from ..schema import DOC_TYPES, FIELD_LABELS, REL_TYPES
from .common import TransformerBlock

NUM_PAIR_FEATS = 8


def pair_features(rel_boxes):
    """rel_boxes: (B, N, 4) boxes in median-text-height units -> (B, N, N, 8) pairwise geometry."""
    x0, y0, x1, y1 = rel_boxes.unbind(-1)
    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
    h = (y1 - y0).clamp(min=1e-3)
    w = (x1 - x0).clamp(min=1e-3)

    def d(a, b):  # b_j - a_i
        return b[:, None, :] - a[:, :, None]

    v_ov = (torch.minimum(y1[:, :, None], y1[:, None, :]) - torch.maximum(y0[:, :, None], y0[:, None, :])).clamp(min=0)
    h_ov = (torch.minimum(x1[:, :, None], x1[:, None, :]) - torch.maximum(x0[:, :, None], x0[:, None, :])).clamp(min=0)
    v_ov = v_ov / torch.minimum(h[:, :, None], h[:, None, :])
    h_ov = h_ov / torch.minimum(w[:, :, None], w[:, None, :])
    return torch.stack([
        d(cy, cy).clamp(-20, 20) / 20,
        d(cx, cx).clamp(-40, 40) / 40,
        v_ov,
        h_ov,
        d(x1, x0).clamp(-40, 40) / 40,
        (v_ov > 0.5).to(rel_boxes.dtype),
        d(x1, x1).abs().clamp(max=20) / 20,
        d(x0, x0).abs().clamp(max=20) / 20,
    ], dim=-1)


class CharEncoder(nn.Module):
    def __init__(self, d, emb=48):
        super().__init__()
        self.emb = nn.Embedding(EXTRACTOR_CLASSES, emb, padding_idx=0)
        self.conv1 = nn.Conv1d(emb, d, 3, padding=1)
        self.conv2 = nn.Conv1d(d, d, 3, padding=1, groups=4)
        self.act = nn.GELU()

    def forward(self, ids):
        """ids: (B, N, L) -> (B, N, d) via masked max-pool over characters."""
        B, N, L = ids.shape
        x = self.emb(ids.reshape(B * N, L)).transpose(1, 2)
        x = self.act(self.conv2(self.act(self.conv1(x))))
        mask = (ids.reshape(B * N, 1, L) == 0)
        x = x.masked_fill(mask, -1e4).amax(-1)
        empty = mask.all(-1)
        x = torch.where(empty, torch.zeros_like(x), x)
        return x.reshape(B, N, -1)


class Extractor(nn.Module):
    def __init__(self, d=112, layers=3, heads=4, ff=256, rel_dim=48, dropout=0.1):
        super().__init__()
        self.heads = heads
        self.chars = CharEncoder(d)
        self.feat_proj = nn.Sequential(nn.Linear(NUM_FEATS, d), nn.GELU(), nn.Linear(d, d))
        self.ln_in = nn.LayerNorm(d)
        self.attn_bias = nn.Sequential(nn.Linear(NUM_PAIR_FEATS, 32), nn.GELU(), nn.Linear(32, heads))
        self.blocks = nn.ModuleList(TransformerBlock(d, heads, ff, dropout) for _ in range(layers))
        self.ln = nn.LayerNorm(d)
        self.field_head = nn.Linear(d, len(FIELD_LABELS))
        self.rel_q = nn.Linear(d, rel_dim * len(REL_TYPES))
        self.rel_k = nn.Linear(d, rel_dim * len(REL_TYPES))
        self.rel_bias = nn.Sequential(nn.Linear(NUM_PAIR_FEATS, 32), nn.GELU(), nn.Linear(32, len(REL_TYPES)))
        self.doc_pool = nn.Linear(d, 1)
        self.doc_head = nn.Linear(d, len(DOC_TYPES))
        self.rel_dim = rel_dim

    def forward(self, char_ids, feats, rel_boxes, valid):
        """char_ids (B,N,L) int64, feats (B,N,F), rel_boxes (B,N,4), valid (B,N) float 1/0.

        Returns field_logits (B,N,K), rel_logits (B,R,N,N), doc_logits (B,D).
        """
        B, N, _ = feats.shape
        x = self.ln_in(self.chars(char_ids) + self.feat_proj(feats))
        pf = pair_features(rel_boxes)                                     # (B,N,N,P)
        key_pad = (1.0 - valid)[:, None, None, :] * -1e4                  # (B,1,1,N)
        bias = self.attn_bias(pf).permute(0, 3, 1, 2) + key_pad           # (B,H,N,N)
        for blk in self.blocks:
            x = blk(x, bias)
        x = self.ln(x)

        field_logits = self.field_head(x)
        R, r = len(REL_TYPES), self.rel_dim
        q = self.rel_q(x).reshape(B, N, R, r).permute(0, 2, 1, 3)
        k = self.rel_k(x).reshape(B, N, R, r).permute(0, 2, 1, 3)
        rel_logits = (q @ k.transpose(-2, -1)) / (r ** 0.5) + self.rel_bias(pf).permute(0, 3, 1, 2)

        w = self.doc_pool(x).squeeze(-1) + (1.0 - valid) * -1e4
        pooled = (w.softmax(-1)[..., None] * x).sum(1)
        return field_logits, rel_logits, self.doc_head(pooled)