inferencelayer 0.2.1

//! TableFormer (docling-models v2.3.0, `accurate`) — CPU, pure Rust, stage-gated vs torch.
//!
//! The neural half of Docling's table-structure stage: `TFPredictor` → `TableModel04_rs.predict`
//! reimplemented on the engine per `MODEL_CONTRACT.md` §2. Ported line-for-line from
//! `docling_ibm_models/tableformer/{data_management/tf_predictor.py,models/table04_rs/*}` (3.10.3)
//! against the checkpoint pinned in `models/docling/tableformer-accurate/PINS.json`
//! (`docling-project/docling-models @ v2.3.0`, `accurate`). Every stage is oracle-gated in
//! `tests/tableformer_parity.rs`.
//!
//! Shape of the thing (all fixed by `tm_config.json`):
//! * prep (`cv_resize`): `resize_img` cv2 INTER_AREA → height 1024 (u8), bbox×sf then crop with
//!   Python `round` (HALF-TO-EVEN), `_prepare_image` = normalize on u8 `(x−255·mean)/std` → cv2
//!   INTER_LINEAR 448×448 (f32) → the `(C,W,H)` transpose quirk → `/255`. The whole chain sits
//!   inside the bit-exact gate (crop byte-exact, prepared ≤1e-5).
//! * encoder (`Encoder04`): torchvision resnet18 truncated after layer3 (`children()[:-3]`) +
//!   AdaptiveAvgPool2d(28) [identity at 28×28] → `[28,28,256]` NHWC. BN folded at load.
//! * tag transformer: `_input_filter` (2 resnet BasicBlocks 256→512) → 784 tokens ×512; a
//!   6-layer post-norm ReLU encoder run ONCE (no positional encoding); a 6-layer post-norm
//!   decoder (sinusoidal PE on the tag embedding, per-layer activation cache == a KV cache) with
//!   per-step greedy argmax over the 13-tag OTSL vocab; structure-error-correction rules verbatim
//!   INCLUDING the upstream never-incremented `line_num` bug (so the xcel→lcel rule is always on).
//! * bbox head (`BBoxDecoder`): its OWN `_input_filter`, CellAttention (softmax over 784) + sigmoid
//!   gate → class head + 3-layer sigmoid-cxcywh MLP, then the `mergebboxes` horizontal-span merge.

use anyhow::{Context, Result};
use std::path::Path;

use crate::conv2d::{Conv2d, adaptive_avg_pool2d, fold_bn, max_pool2d};
use crate::cpu_gemm::{PackedWeight, gemm_packed};
use crate::cv_resize::{inter_area_u8, inter_linear_f32};
use crate::weights::LazySt;

// ── dims / vocab (tm_config.json) ─────────────────────────────────────────────────────────────────
const D: usize = 512; // hidden_dim / tag_decoder_dim
const HEADS: usize = 8;
const HDIM: usize = D / HEADS; // 64
const GRID: usize = 28; // enc_image_size
const TOKENS: usize = GRID * GRID; // 784
const VOCAB: usize = 13;
const ENC_LAYERS: usize = 6;
const DEC_LAYERS: usize = 6;
const RESIZED: usize = 448;
const MAX_STEPS: usize = 1024;
const LN_EPS: f32 = 1e-5;
const BN_EPS: f32 = 1e-5;

// OTSL tag ids (tm_config `word_map_tag`).
const START: u32 = 2;
const END: u32 = 3;
const ECEL: u32 = 4;
const FCEL: u32 = 5;
const LCEL: u32 = 6;
const UCEL: u32 = 7;
const XCEL: u32 = 8;
const NL: u32 = 9;
const CHED: u32 = 10;
const RHED: u32 = 11;
const SROW: u32 = 12;

const TAG_STR: [&str; VOCAB] = [
    "<pad>", "<unk>", "<start>", "<end>", "ecel", "fcel", "lcel", "ucel", "xcel", "nl", "ched",
    "rhed", "srow",
];

// tm_config image_normalization mean/std, kept f64 — torch normalizes in f64 (numpy promotes the
// python-list mean/std), then casts to f32 only at the very end. We match that: normalize in f64.
const MEAN: [f64; 3] = [0.94247851, 0.94254675, 0.94292611];
const STD: [f64; 3] = [0.17910956, 0.17940403, 0.17931663];

/// `MODEL_CONTRACT.md` §2 output: greedy OTSL tags + the (merged) bbox head outputs in crop space.
#[derive(Clone, Debug)]
pub struct TableStructure {
    /// decoded tag strings, greedy, post error-correction, WITHOUT `<start>`/`<end>`
    /// (== `_get_html_tags`).
    pub otsl_seq: Vec<String>,
    /// argmax of the bbox-decoder class head, aligned with `bboxes` (post-mergebboxes).
    pub cell_classes: Vec<u32>,
    /// cxcywh, NORMALIZED [0,1] in CROP space, post-mergebboxes.
    pub bboxes: Vec<[f32; 4]>,
}

/// Every intermediate the parity gate scores, mirroring the oracle's stage keys.
pub struct Stages {
    /// u8 crop after `resize_img` + round-crop (`_prepare_image` input) — `[ch, cw, 3]`.
    pub crop: Vec<u8>,
    pub crop_h: usize,
    pub crop_w: usize,
    /// `[1,3,448,448]` prepared tensor (`_prepare_image` output).
    pub prepared: Vec<f32>,
    /// encoder output `[28,28,256]` NHWC.
    pub encoder_out: Vec<f32>,
    /// tag-transformer `_input_filter` output `[28,28,512]` NHWC (oracle stores it NCHW).
    pub input_filter_out: Vec<f32>,
    /// full greedy tag id sequence incl. `<start>`/`<end>` (== `decoded_tags`).
    pub seq: Vec<u32>,
    /// bbox-decoder class head, post-merge `[n,3]`.
    pub outputs_class: Vec<[f32; 3]>,
    /// bbox-decoder coords (sigmoid cxcywh), post-merge `[n,4]`.
    pub outputs_coord: Vec<[f32; 4]>,
}

// ── small nn pieces (private copies, deliberately uncoupled from rtdetr) ──────────────────────────

/// Dense layer, torch `[n,k]` weight prepacked for the GEMM.
struct Lin {
    w: PackedWeight,
    b: Vec<f32>,
}

impl Lin {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let shape = st.shape(&format!("{prefix}.weight"))?;
        let (n, k) = (shape[0], shape[1]);
        let b = st.tensor_f32(&format!("{prefix}.bias"))?;
        Ok(Self::from_parts(w, n, k, b))
    }

    fn from_parts(w: Vec<f32>, n: usize, k: usize, b: Vec<f32>) -> Self {
        Self {
            w: PackedWeight::new(&w, n, k),
            b,
        }
    }

    fn forward(&self, x: &[f32], m: usize) -> Vec<f32> {
        let mut out = vec![0f32; m * self.w.n()];
        gemm_packed(&mut out, x, &self.w, m, Some(&self.b));
        out
    }
}

struct LayerNorm {
    g: Vec<f32>,
    b: Vec<f32>,
}

impl LayerNorm {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            g: st.tensor_f32(&format!("{prefix}.weight"))?,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
        })
    }

    fn forward_inplace(&self, x: &mut [f32]) {
        let d = self.g.len();
        for row in x.chunks_mut(d) {
            let mean = row.iter().sum::<f32>() / d as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
            let inv = 1.0 / (var + LN_EPS).sqrt();
            for (i, v) in row.iter_mut().enumerate() {
                *v = (*v - mean) * inv * self.g[i] + self.b[i];
            }
        }
    }
}

fn relu_inplace(x: &mut [f32]) {
    x.iter_mut().for_each(|v| *v = v.max(0.0));
}

fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

/// Plain multi-head attention `[nq, 512] × [nk, 512] → [nq, 512]`, 8 heads × 64, 1/√64 scaling.
/// `q`,`k`,`v` are ALREADY projected. Softmax with the standard max-subtract.
///
/// Perf shape (the port's measured hotspot — 765 of 1115 ms at 784 tokens): rayon over QUERY
/// rows (each output row's arithmetic order is unchanged, so results are bit-identical to the
/// serial loop) and `simd::dot` for the QKᵀ reductions — the scalar `.sum()` cannot vectorize
/// (sequential float semantics). The NEON dot reassociates the 64-term reduction; the
/// tag-sequence exactness gate in tableformer_parity.rs is the arbiter of that tolerance.
fn mha512(q: &[f32], k: &[f32], v: &[f32], nq: usize, nk: usize) -> Vec<f32> {
    use rayon::prelude::*;
    let scale = 1.0 / (HDIM as f32).sqrt();
    let mut out = vec![0f32; nq * D];
    let row = |i: usize, orow: &mut [f32]| {
        let mut scores = vec![0f32; nk];
        for h in 0..HEADS {
            let off = h * HDIM;
            let qi = &q[i * D + off..i * D + off + HDIM];
            let mut maxv = f32::NEG_INFINITY;
            for (j, s) in scores.iter_mut().enumerate() {
                let kj = &k[j * D + off..j * D + off + HDIM];
                *s = crate::simd::dot(qi, kj) * scale;
                maxv = maxv.max(*s);
            }
            let mut denom = 0f32;
            for s in scores.iter_mut() {
                *s = (*s - maxv).exp();
                denom += *s;
            }
            let inv = 1.0 / denom;
            let oi = &mut orow[off..off + HDIM];
            for (j, s) in scores.iter().enumerate() {
                let wj = s * inv;
                let vj = &v[j * D + off..j * D + off + HDIM];
                for (o, vv) in oi.iter_mut().zip(vj) {
                    *o += wj * vv;
                }
            }
        }
    };
    if nq == 1 {
        // the decode loop's per-step calls — rayon dispatch would be pure overhead
        row(0, &mut out);
        return out;
    }
    out.par_chunks_mut(D)
        .enumerate()
        .for_each(|(i, orow)| row(i, orow));
    out
}

/// nn.MultiheadAttention packed in-proj (`in_proj_weight [1536,512]`, `in_proj_bias [1536]`)
/// split into Wq/Wk/Wv + the out projection.
struct Mha {
    wq: Lin,
    wk: Lin,
    wv: Lin,
    out: Lin,
}

impl Mha {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.in_proj_weight"))?; // [3D, D]
        let b = st.tensor_f32(&format!("{prefix}.in_proj_bias"))?; // [3D]
        let split = |i: usize| -> Lin {
            Lin::from_parts(
                w[i * D * D..(i + 1) * D * D].to_vec(),
                D,
                D,
                b[i * D..(i + 1) * D].to_vec(),
            )
        };
        Ok(Self {
            wq: split(0),
            wk: split(1),
            wv: split(2),
            out: Lin::load(st, &format!("{prefix}.out_proj"))?,
        })
    }
}

// ── resnet building blocks (torchvision BasicBlock, BN folded at load) ────────────────────────────

/// Conv + folded eval-BN (+ optional ReLU). `pad = (k-1)/2` (torchvision `conv3x3`/`conv1x1`).
struct ConvBn {
    conv: Conv2d,
    relu: bool,
}

impl ConvBn {
    fn load(st: &LazySt, conv: &str, bn: &str, stride: usize, relu: bool) -> Result<Self> {
        let w = st.tensor_f32(&format!("{conv}.weight"))?;
        let shape = st.shape(&format!("{conv}.weight"))?.to_vec();
        let (oc, ic, kh, kw) = (shape[0], shape[1], shape[2], shape[3]);
        let gamma = st.tensor_f32(&format!("{bn}.weight"))?;
        let beta = st.tensor_f32(&format!("{bn}.bias"))?;
        let mean = st.tensor_f32(&format!("{bn}.running_mean"))?;
        let var = st.tensor_f32(&format!("{bn}.running_var"))?;
        let (wf, bf) = fold_bn(&w, None, oc, &gamma, &beta, &mean, &var, BN_EPS);
        let pad = (kh - 1) / 2;
        Ok(Self {
            conv: Conv2d::from_torch(&wf, Some(&bf), oc, ic, kh, kw, stride, pad),
            relu,
        })
    }

    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (mut o, oh, ow) = self.conv.forward(x, h, w);
        if self.relu {
            relu_inplace(&mut o);
        }
        (o, oh, ow)
    }
}

/// torchvision `BasicBlock`: conv1(relu)→conv2→(+ downsample(x))→relu. Stride lives in conv1 (and
/// the downsample) — matches `conv3x3(inplanes,planes,stride)`.
struct BasicBlock {
    conv1: ConvBn,
    conv2: ConvBn,
    downsample: Option<ConvBn>,
}

impl BasicBlock {
    fn load(st: &LazySt, base: &str, stride: usize, has_ds: bool) -> Result<Self> {
        let downsample = if has_ds {
            Some(ConvBn::load(
                st,
                &format!("{base}.downsample.0"),
                &format!("{base}.downsample.1"),
                stride,
                false,
            )?)
        } else {
            None
        };
        Ok(Self {
            conv1: ConvBn::load(
                st,
                &format!("{base}.conv1"),
                &format!("{base}.bn1"),
                stride,
                true,
            )?,
            conv2: ConvBn::load(
                st,
                &format!("{base}.conv2"),
                &format!("{base}.bn2"),
                1,
                false,
            )?,
            downsample,
        })
    }

    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (a, h1, w1) = self.conv1.forward(x, h, w);
        let (mut a, h2, w2) = self.conv2.forward(&a, h1, w1);
        let owned = self.downsample.as_ref().map(|ds| ds.forward(x, h, w).0);
        let res: &[f32] = owned.as_deref().unwrap_or(x);
        debug_assert_eq!(a.len(), res.len());
        for (v, r) in a.iter_mut().zip(res) {
            *v = (*v + r).max(0.0);
        }
        (a, h2, w2)
    }
}

/// `u.resnet_block(stride=1)` — the two `_input_filter` BasicBlocks (256→512 with 1×1 downsample,
/// then 512→512). NHWC in, NHWC out.
struct InputFilter {
    b0: BasicBlock,
    b1: BasicBlock,
}

impl InputFilter {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        Ok(Self {
            b0: BasicBlock::load(st, &format!("{base}.0"), 1, true)?,
            b1: BasicBlock::load(st, &format!("{base}.1"), 1, false)?,
        })
    }

    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (x, h, w) = self.b0.forward(x, h, w);
        self.b1.forward(&x, h, w)
    }
}

/// `Encoder04`: torchvision resnet18 truncated after layer3 + AdaptiveAvgPool2d(28).
struct Encoder {
    stem_conv: ConvBn, // conv1 7×7 s2 + bn1 + relu
    layer1: [BasicBlock; 2],
    layer2: [BasicBlock; 2],
    layer3: [BasicBlock; 2],
}

impl Encoder {
    fn load(st: &LazySt) -> Result<Self> {
        let p = "_encoder._resnet";
        let block = |idx: usize, sub: usize, stride: usize, has_ds: bool| {
            BasicBlock::load(st, &format!("{p}.{idx}.{sub}"), stride, has_ds)
        };
        Ok(Self {
            stem_conv: ConvBn::load(st, &format!("{p}.0"), &format!("{p}.1"), 2, true)?,
            layer1: [block(4, 0, 1, false)?, block(4, 1, 1, false)?],
            layer2: [block(5, 0, 2, true)?, block(5, 1, 1, false)?],
            layer3: [block(6, 0, 2, true)?, block(6, 1, 1, false)?],
        })
    }

    /// `[448,448,3]` NHWC → `[28,28,256]` NHWC.
    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (x, h, w) = self.stem_conv.forward(x, h, w); // → 224×224×64
        let (mut x, mut h, mut w) = max_pool2d(&x, h, w, 64, 3, 2, 1); // → 112×112×64
        for stage in [&self.layer1, &self.layer2, &self.layer3] {
            for b in stage {
                let (nx, nh, nw) = b.forward(&x, h, w);
                x = nx;
                h = nh;
                w = nw;
            }
        }
        // AdaptiveAvgPool2d(28) is identity at 28×28; keep the call only if it ever isn't.
        if (h, w) != (GRID, GRID) {
            x = adaptive_avg_pool2d(&x, h, w, 256, GRID, GRID);
        }
        (x, GRID, GRID)
    }
}

// ── tag transformer ───────────────────────────────────────────────────────────────────────────────

/// post-norm encoder layer (ReLU FFN, no positional encoding).
struct EncLayer {
    attn: Mha,
    norm1: LayerNorm,
    lin1: Lin,
    lin2: Lin,
    norm2: LayerNorm,
}

impl EncLayer {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        Ok(Self {
            attn: Mha::load(st, &format!("{base}.self_attn"))?,
            norm1: LayerNorm::load(st, &format!("{base}.norm1"))?,
            lin1: Lin::load(st, &format!("{base}.linear1"))?,
            lin2: Lin::load(st, &format!("{base}.linear2"))?,
            norm2: LayerNorm::load(st, &format!("{base}.norm2"))?,
        })
    }

    fn forward(&self, x: &[f32], n: usize) -> Vec<f32> {
        let q = self.attn.wq.forward(x, n);
        let k = self.attn.wk.forward(x, n);
        let v = self.attn.wv.forward(x, n);
        let a = self.attn.out.forward(&mha512(&q, &k, &v, n, n), n);
        let mut h: Vec<f32> = x.iter().zip(&a).map(|(p, q)| p + q).collect();
        self.norm1.forward_inplace(&mut h);
        let mut f = self.lin1.forward(&h, n);
        relu_inplace(&mut f);
        let f = self.lin2.forward(&f, n);
        let mut o: Vec<f32> = h.iter().zip(&f).map(|(p, q)| p + q).collect();
        self.norm2.forward_inplace(&mut o);
        o
    }
}

/// `TMTransformerDecoderLayer`: post-norm, ReLU FFN. Self-attn over the cached layer-input history
/// (query = the new token only), then cross-attn to the tag-encoder memory.
struct DecLayer {
    self_attn: Mha,
    norm1: LayerNorm,
    cross_attn: Mha,
    norm2: LayerNorm,
    lin1: Lin,
    lin2: Lin,
    norm3: LayerNorm,
}

impl DecLayer {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        Ok(Self {
            self_attn: Mha::load(st, &format!("{base}.self_attn"))?,
            norm1: LayerNorm::load(st, &format!("{base}.norm1"))?,
            cross_attn: Mha::load(st, &format!("{base}.multihead_attn"))?,
            norm2: LayerNorm::load(st, &format!("{base}.norm2"))?,
            lin1: Lin::load(st, &format!("{base}.linear1"))?,
            lin2: Lin::load(st, &format!("{base}.linear2"))?,
            norm3: LayerNorm::load(st, &format!("{base}.norm3"))?,
        })
    }

    /// One decode step. `hist_in` is this layer's full input history `[(t+1)·512]` (the new token
    /// is the last row); `mem_k`/`mem_v` are the precomputed cross-attn key/value projections of
    /// the encoder memory. Returns the layer output for the new token `[512]`.
    /// One decode step with an incremental self-attn K/V cache: the K/V projections are
    /// row-independent (each output row is a dot-product chain over the 512 inputs of THAT row
    /// alone), so projecting only the new position and appending yields bit-identical K/V to
    /// re-projecting the whole history every step — the full-tag-sequence parity gate arbitrates.
    /// Dropping the O(T²) re-projection is what makes monster tables (T ~ 1000) affordable.
    fn step(
        &self,
        cur: &[f32],
        t: usize,
        k_cache: &mut Vec<f32>,
        v_cache: &mut Vec<f32>,
        mem_k: &[f32],
        mem_v: &[f32],
    ) -> Vec<f32> {
        // self-attn: query = last token, keys/values = cached history + this position.
        let q = self.self_attn.wq.forward(cur, 1);
        k_cache.extend_from_slice(&self.self_attn.wk.forward(cur, 1));
        v_cache.extend_from_slice(&self.self_attn.wv.forward(cur, 1));
        let a = self
            .self_attn
            .out
            .forward(&mha512(&q, k_cache, v_cache, 1, t + 1), 1);
        let mut h1: Vec<f32> = cur.iter().zip(&a).map(|(p, q)| p + q).collect();
        self.norm1.forward_inplace(&mut h1);
        // cross-attn: query = h1, keys/values = memory (K,V precomputed once).
        let q2 = self.cross_attn.wq.forward(&h1, 1);
        let a2 = self
            .cross_attn
            .out
            .forward(&mha512(&q2, mem_k, mem_v, 1, TOKENS), 1);
        let mut h2: Vec<f32> = h1.iter().zip(&a2).map(|(p, q)| p + q).collect();
        self.norm2.forward_inplace(&mut h2);
        // FFN
        let mut f = self.lin1.forward(&h2, 1);
        relu_inplace(&mut f);
        let f = self.lin2.forward(&f, 1);
        let mut h3: Vec<f32> = h2.iter().zip(&f).map(|(p, q)| p + q).collect();
        self.norm3.forward_inplace(&mut h3);
        h3
    }
}

struct TagTransformer {
    input_filter: InputFilter,
    embedding: Vec<f32>,    // [13,512]
    pe: Vec<f32>,           // [1024,512] (buffer, batch dim dropped)
    encoder: Vec<EncLayer>, // 6
    decoder: Vec<DecLayer>, // 6
    fc: Lin,                // [13,512]
}

impl TagTransformer {
    fn load(st: &LazySt) -> Result<Self> {
        let p = "_tag_transformer";
        let mut encoder = Vec::with_capacity(ENC_LAYERS);
        for i in 0..ENC_LAYERS {
            encoder.push(EncLayer::load(st, &format!("{p}._encoder.layers.{i}"))?);
        }
        let mut decoder = Vec::with_capacity(DEC_LAYERS);
        for i in 0..DEC_LAYERS {
            decoder.push(DecLayer::load(st, &format!("{p}._decoder.layers.{i}"))?);
        }
        // positional_encoding.pe is [1024,1,512]; drop the batch dim.
        let pe = st.tensor_f32(&format!("{p}._positional_encoding.pe"))?;
        Ok(Self {
            input_filter: InputFilter::load(st, &format!("{p}._input_filter"))?,
            embedding: st.tensor_f32(&format!("{p}._embedding.weight"))?,
            pe,
            encoder,
            decoder,
            fc: Lin::load(st, &format!("{p}._fc"))?,
        })
    }

    /// embed a tag + add its positional encoding at position `t`.
    fn embed(&self, tag: u32, t: usize) -> Vec<f32> {
        let e = &self.embedding[tag as usize * D..(tag as usize + 1) * D];
        let p = &self.pe[t * D..(t + 1) * D];
        e.iter().zip(p).map(|(a, b)| a + b).collect()
    }
}

// ── bbox decoder ──────────────────────────────────────────────────────────────────────────────────

/// `CellAttention` (3 projections + full_att) + gate + heads.
struct BBoxDecoder {
    input_filter: InputFilter,
    encoder_att: Lin,
    tag_decoder_att: Lin,
    language_att: Lin,
    full_att: Lin,
    init_h: Lin,
    f_beta: Lin,
    class_embed: Lin,     // [3,512]
    bbox_embed: Vec<Lin>, // MLP 512→256→256→4
}

impl BBoxDecoder {
    fn load(st: &LazySt) -> Result<Self> {
        let p = "_bbox_decoder";
        let mut bbox_embed = Vec::with_capacity(3);
        for i in 0..3 {
            bbox_embed.push(Lin::load(st, &format!("{p}._bbox_embed.layers.{i}"))?);
        }
        Ok(Self {
            input_filter: InputFilter::load(st, &format!("{p}._input_filter"))?,
            encoder_att: Lin::load(st, &format!("{p}._attention._encoder_att"))?,
            tag_decoder_att: Lin::load(st, &format!("{p}._attention._tag_decoder_att"))?,
            language_att: Lin::load(st, &format!("{p}._attention._language_att"))?,
            full_att: Lin::load(st, &format!("{p}._attention._full_att"))?,
            init_h: Lin::load(st, &format!("{p}._init_h"))?,
            f_beta: Lin::load(st, &format!("{p}._f_beta"))?,
            class_embed: Lin::load(st, &format!("{p}._class_embed"))?,
            bbox_embed,
        })
    }

    /// `inference(enc_out, tag_H)` — `enc_out` is the RAW encoder `[28,28,256]` NHWC; `tag_h` are
    /// the captured per-cell decoder outputs `[512]`. Returns `(classes [n,3], coords [n,4])`.
    fn inference(&self, enc_out: &[f32], tag_h: &[Vec<f32>]) -> (Vec<[f32; 3]>, Vec<[f32; 4]>) {
        // this decoder's OWN _input_filter → [28,28,512] → 784 tokens.
        let (enc, _, _) = self.input_filter.forward(enc_out, GRID, GRID);

        // pieces independent of the cell (h depends only on the mean encoding).
        let mut mean_enc = vec![0f32; D];
        for tok in 0..TOKENS {
            for c in 0..D {
                mean_enc[c] += enc[tok * D + c];
            }
        }
        mean_enc.iter_mut().for_each(|v| *v /= TOKENS as f32);
        let h = self.init_h.forward(&mean_enc, 1); // [512]
        let att1 = self.encoder_att.forward(&enc, TOKENS); // [784,512]
        let att3 = self.language_att.forward(&h, 1); // [512]
        // base[tok] = att1[tok] + att3
        let mut base = att1;
        for tok in 0..TOKENS {
            for c in 0..D {
                base[tok * D + c] += att3[c];
            }
        }
        let gate: Vec<f32> = self
            .f_beta
            .forward(&h, 1)
            .iter()
            .map(|&v| sigmoid(v))
            .collect();

        let mut classes = Vec::with_capacity(tag_h.len());
        let mut coords = Vec::with_capacity(tag_h.len());
        for cell in tag_h {
            let att2 = self.tag_decoder_att.forward(cell, 1); // [512]
            // pre[tok] = relu(base[tok] + att2); att[tok] = full_att(pre[tok])
            let mut pre = vec![0f32; TOKENS * D];
            for tok in 0..TOKENS {
                for c in 0..D {
                    pre[tok * D + c] = (base[tok * D + c] + att2[c]).max(0.0);
                }
            }
            let att = self.full_att.forward(&pre, TOKENS); // [784,1]
            // softmax over 784
            let mut maxv = f32::NEG_INFINITY;
            for &a in &att {
                maxv = maxv.max(a);
            }
            let mut denom = 0f32;
            let mut alpha = vec![0f32; TOKENS];
            for (a, al) in att.iter().zip(alpha.iter_mut()) {
                *al = (a - maxv).exp();
                denom += *al;
            }
            let inv = 1.0 / denom;
            // awe = Σ enc[tok]*alpha[tok]; then gate⊙awe; then ⊙h
            let mut awe = vec![0f32; D];
            for tok in 0..TOKENS {
                let a = alpha[tok] * inv;
                for c in 0..D {
                    awe[c] += enc[tok * D + c] * a;
                }
            }
            let h2: Vec<f32> = (0..D).map(|c| gate[c] * awe[c] * h[c]).collect();

            // class head + 3-layer sigmoid-cxcywh MLP
            let cls = self.class_embed.forward(&h2, 1);
            classes.push([cls[0], cls[1], cls[2]]);
            let mut b = h2.clone();
            for (i, l) in self.bbox_embed.iter().enumerate() {
                b = l.forward(&b, 1);
                if i + 1 < self.bbox_embed.len() {
                    relu_inplace(&mut b);
                }
            }
            coords.push([sigmoid(b[0]), sigmoid(b[1]), sigmoid(b[2]), sigmoid(b[3])]);
        }
        (classes, coords)
    }
}

// ── the model ───────────────────────────────────────────────────────────────────────────────────

pub struct TableFormer {
    encoder: Encoder,
    tag: TagTransformer,
    bbox: BBoxDecoder,
}

impl TableFormer {
    /// Load from `models/docling/tableformer-accurate` (the pinned `tableformer_accurate.safetensors`).
    pub fn load(dir: &Path) -> Result<Self> {
        let path = dir.join("tableformer_accurate.safetensors");
        let bytes = std::fs::read(&path)
            .with_context(|| format!("read tableformer checkpoint {}", path.display()))?;
        let st = LazySt::from_bytes(vec![bytes]).context("parse tableformer safetensors")?;
        Ok(Self {
            encoder: Encoder::load(&st)?,
            tag: TagTransformer::load(&st)?,
            bbox: BBoxDecoder::load(&st)?,
        })
    }

    /// `MODEL_CONTRACT.md` §2 surface: page RGB8 (rendered at scale 2.0) + a table bbox in that
    /// image's pixel space → OTSL tags + merged bbox-head outputs (crop space).
    pub fn recognize(
        &self,
        page_rgb: &[u8],
        page_w: usize,
        page_h: usize,
        table_bbox: [f32; 4],
    ) -> TableStructure {
        let s = self.recognize_stages(page_rgb, page_w, page_h, table_bbox);
        // otsl_seq == _get_html_tags: drop <start>/<end>, map ids to strings.
        let otsl_seq = s.seq[1..s.seq.len().saturating_sub(1)]
            .iter()
            .map(|&t| TAG_STR[t as usize].to_string())
            .collect();
        let cell_classes = s.outputs_class.iter().map(|c| argmax3(c) as u32).collect();
        TableStructure {
            otsl_seq,
            cell_classes,
            bboxes: s.outputs_coord,
        }
    }

    /// Full forward with every oracle-gated stage exposed. `OSFKB_TABLEFORMER_TRACE=1` prints
    /// per-stage wall time (house profiling convention).
    pub fn recognize_stages(
        &self,
        page_rgb: &[u8],
        page_w: usize,
        page_h: usize,
        table_bbox: [f32; 4],
    ) -> Stages {
        let trace = std::env::var("OSFKB_TABLEFORMER_TRACE").is_ok();
        let mut t = std::time::Instant::now();
        let mut lap = |name: &str| {
            if trace {
                eprintln!(
                    "tableformer {name}: {:.1} ms",
                    t.elapsed().as_secs_f64() * 1e3
                );
            }
            t = std::time::Instant::now();
        };

        // ── prep: resize_img (height 1024, INTER_AREA) → scale bbox → round-crop ──
        // geometry in f64 (Python floats), so int()-truncation and round()-to-even match exactly.
        let sf = 1024.0f64 / page_h as f64;
        let rw = (page_w as f64 * sf) as usize; // width TRUNCATED by int()
        let resized = inter_area_u8(page_rgb, page_h, page_w, 3, 1024, rw);
        let sb = [
            table_bbox[0] as f64 * sf,
            table_bbox[1] as f64 * sf,
            table_bbox[2] as f64 * sf,
            table_bbox[3] as f64 * sf,
        ];
        let (l, t0, r, b) = (
            py_round(sb[0]),
            py_round(sb[1]),
            py_round(sb[2]),
            py_round(sb[3]),
        );
        let (cw, ch) = (r - l, b - t0);
        let mut crop = vec![0u8; ch * cw * 3];
        for y in 0..ch {
            let src = ((t0 + y) * rw + l) * 3;
            crop[y * cw * 3..(y + 1) * cw * 3].copy_from_slice(&resized[src..src + cw * 3]);
        }
        lap("prep-crop");

        // ── _prepare_image: normalize (u8) → INTER_LINEAR 448 → (C,W,H) transpose → /255 ──
        let mut norm = vec![0f32; ch * cw * 3];
        for i in 0..ch * cw {
            for c in 0..3 {
                norm[i * 3 + c] = ((crop[i * 3 + c] as f64 - 255.0 * MEAN[c]) / STD[c]) as f32;
            }
        }
        let resized_lin = inter_linear_f32(&norm, ch, cw, 3, RESIZED, RESIZED); // [448,448,3] HWC
        // prepared[0,c,a,b] = resized_lin[H=b, W=a, C=c] / 255  (the (C,W,H) transpose quirk)
        let mut prepared = vec![0f32; 3 * RESIZED * RESIZED];
        for c in 0..3 {
            for a in 0..RESIZED {
                for bb in 0..RESIZED {
                    prepared[(c * RESIZED + a) * RESIZED + bb] =
                        resized_lin[(bb * RESIZED + a) * 3 + c] / 255.0;
                }
            }
        }
        lap("prepare-image");

        // ── encoder + input_filter ──
        // torch runs the encoder on `prepared` (NCHW, WITH the (C,W,H) transpose quirk baked in),
        // so the engine's NHWC encoder gets the NCHW→NHWC view of `prepared` — NOT the plain HWC
        // resize (transposing H/W changes the image; the quirk is preserved, not fixed).
        let mut enc_in = vec![0f32; RESIZED * RESIZED * 3];
        for c in 0..3 {
            for a in 0..RESIZED {
                for bb in 0..RESIZED {
                    enc_in[(a * RESIZED + bb) * 3 + c] = prepared[(c * RESIZED + a) * RESIZED + bb];
                }
            }
        }
        let (encoder_out, _, _) = self.encoder.forward(&enc_in, RESIZED, RESIZED); // [28,28,256]
        lap("encoder");
        let (input_filter_out, _, _) = self.tag.input_filter.forward(&encoder_out, GRID, GRID); // [28,28,512]
        lap("input-filter");

        // ── tag encoder (once) ──
        let memory = {
            let mut m = input_filter_out.clone(); // [784,512] token-major
            for layer in &self.tag.encoder {
                m = layer.forward(&m, TOKENS);
            }
            m
        };
        // precompute cross-attn K,V of memory, per decoder layer.
        let mem_kv: Vec<(Vec<f32>, Vec<f32>)> = self
            .tag
            .decoder
            .iter()
            .map(|l| {
                (
                    l.cross_attn.wk.forward(&memory, TOKENS),
                    l.cross_attn.wv.forward(&memory, TOKENS),
                )
            })
            .collect();
        lap("tag-encoder");

        // ── greedy decode with structure error-correction + bbox bookkeeping ──
        let (seq, tag_h, merge) = self.decode(&mem_kv);
        lap("decode");

        // ── bbox head + mergebboxes ──
        let (cls, coord) = self.bbox.inference(&encoder_out, &tag_h);
        let (outputs_class, outputs_coord) = merge_bboxes(&cls, &coord, &merge);
        lap("bbox+merge");

        Stages {
            crop,
            crop_h: ch,
            crop_w: cw,
            prepared,
            encoder_out,
            input_filter_out,
            seq,
            outputs_class,
            outputs_coord,
        }
    }

    /// The greedy loop of `TableModel04_rs.predict`. Returns the tag sequence (incl. start/end),
    /// the captured per-cell decoder outputs, and the horizontal-span merge map.
    fn decode(&self, mem_kv: &[(Vec<f32>, Vec<f32>)]) -> (Vec<u32>, Vec<Vec<f32>>, MergeMap) {
        let dec = &self.tag.decoder;
        // per-layer incremental self-attn K/V caches (layer-0 keys come from the embeddings,
        // layer-l keys from layer-(l-1)'s outputs — same data flow as the full-history version).
        let mut k_caches: Vec<Vec<f32>> = vec![Vec::new(); DEC_LAYERS];
        let mut v_caches: Vec<Vec<f32>> = vec![Vec::new(); DEC_LAYERS];

        let mut seq: Vec<u32> = vec![START];
        let mut tag_h: Vec<Vec<f32>> = Vec::new();

        // error-correction / bookkeeping state (names mirror the reference exactly).
        let mut skip_next_tag = true;
        let mut prev_tag_ucel = false;
        let line_num = 0; // NEVER incremented upstream — the xcel→lcel rule is always on.
        let mut first_lcel = true;
        let mut merge_map = MergeMap::default();
        let mut cur_bbox_ind: i64 = -1;
        let mut bbox_ind: usize = 0;

        let mut steps = 0usize;
        while steps < MAX_STEPS {
            let t = steps;
            let tok = seq[t];
            // run the 6 layers for position t; the final-layer output at position t is decoded_last.
            let mut cur = self.tag.embed(tok, t);
            for l in 0..DEC_LAYERS {
                let (mk, mv) = &mem_kv[l];
                let (kc, vc) = (&mut k_caches[l], &mut v_caches[l]);
                cur = dec[l].step(&cur, t, kc, vc, mk, mv);
            }
            let decoded_last = cur;

            let logits = self.tag.fc.forward(&decoded_last, 1);
            let mut new_tag = argmax(&logits) as u32;

            // STRUCTURE ERROR CORRECTION (verbatim).
            if line_num == 0 && new_tag == XCEL {
                new_tag = LCEL;
            }
            if prev_tag_ucel && new_tag == LCEL {
                new_tag = FCEL;
            }

            if new_tag == END {
                seq.push(END);
                break;
            }

            // ── bbox bookkeeping ──
            if !skip_next_tag && matches!(new_tag, FCEL | ECEL | CHED | RHED | SROW | NL | UCEL) {
                tag_h.push(decoded_last.clone());
                if !first_lcel {
                    merge_map.set(cur_bbox_ind as usize, bbox_ind as i64);
                }
                bbox_ind += 1;
            }
            if new_tag != LCEL {
                first_lcel = true;
            } else if first_lcel {
                tag_h.push(decoded_last.clone());
                first_lcel = false;
                cur_bbox_ind = bbox_ind as i64;
                merge_map.set(bbox_ind, -1);
                bbox_ind += 1;
            }
            skip_next_tag = matches!(new_tag, NL | UCEL | XCEL);
            prev_tag_ucel = new_tag == UCEL;

            seq.push(new_tag);
            steps += 1;
        }
        (seq, tag_h, merge_map)
    }
}

/// horizontal-span merge map (`bboxes_to_merge`): start-index → partner-index (or -1 = unclosed).
#[derive(Default)]
struct MergeMap {
    keys: Vec<usize>,
    vals: Vec<i64>,
}

impl MergeMap {
    fn set(&mut self, k: usize, v: i64) {
        if let Some(pos) = self.keys.iter().position(|&x| x == k) {
            self.vals[pos] = v;
        } else {
            self.keys.push(k);
            self.vals.push(v);
        }
    }
    fn get(&self, k: usize) -> Option<i64> {
        self.keys.iter().position(|&x| x == k).map(|p| self.vals[p])
    }
}

/// `mergebboxes` on the coords the merge map marks, dropping merge partners (`boxes_to_skip`).
/// Python negative-index semantics preserved (`-1` partner → the LAST box).
fn merge_bboxes(
    cls: &[[f32; 3]],
    coord: &[[f32; 4]],
    merge: &MergeMap,
) -> (Vec<[f32; 3]>, Vec<[f32; 4]>) {
    let n = coord.len();
    let mut skip: Vec<i64> = Vec::new();
    let mut out_cls = Vec::new();
    let mut out_coord = Vec::new();
    for box_ind in 0..n {
        let b1 = coord[box_ind];
        let c1 = cls[box_ind];
        if let Some(partner) = merge.get(box_ind) {
            let pidx = if partner < 0 {
                (n as i64 + partner) as usize
            } else {
                partner as usize
            };
            let b2 = coord[pidx];
            skip.push(partner);
            out_coord.push(mergebboxes(&b1, &b2));
            out_cls.push(c1);
        } else if !skip.contains(&(box_ind as i64)) {
            out_coord.push(b1);
            out_cls.push(c1);
        }
    }
    (out_cls, out_coord)
}

/// `TableModel04_rs.mergebboxes` — union of two cxcywh boxes into a cxcywh box.
fn mergebboxes(b1: &[f32; 4], b2: &[f32; 4]) -> [f32; 4] {
    let new_w = (b2[0] + b2[2] / 2.0) - (b1[0] - b1[2] / 2.0);
    let new_h = (b2[1] + b2[3] / 2.0) - (b1[1] - b1[3] / 2.0);
    let new_left = b1[0] - b1[2] / 2.0;
    let new_top = (b2[1] - b2[3] / 2.0).min(b1[1] - b1[3] / 2.0);
    let new_cx = new_left + new_w / 2.0;
    let new_cy = new_top + new_h / 2.0;
    [new_cx, new_cy, new_w, new_h]
}

fn argmax(x: &[f32]) -> usize {
    let mut best = 0;
    for i in 1..x.len() {
        if x[i] > x[best] {
            best = i;
        }
    }
    best
}

fn argmax3(x: &[f32; 3]) -> usize {
    let mut best = 0;
    for i in 1..3 {
        if x[i] > x[best] {
            best = i;
        }
    }
    best
}

/// Python `round` — round half to EVEN (the crop trap; Rust's `f32::round` is half-away-from-zero).
fn py_round(x: f64) -> usize {
    x.round_ties_even().max(0.0) as usize
}