inferencelayer 0.2.1

//! RT-DETRv2 layout detector (docling-layout-heron) — CPU, pure Rust, stage-gated vs transformers.
//!
//! The neural half of Docling's layout stage: `LayoutPredictor.predict()` reimplemented on the
//! engine per `MODEL_CONTRACT.md` — page image in, thresholded/labelled/clamped detections out.
//! Ported line-for-line from `transformers/models/rt_detr_v2/modeling_rt_detr_v2.py` (4.57.3) and
//! `rt_detr/modeling_rt_detr_resnet.py`, against the checkpoint pinned in
//! `tests/fixtures/docling_pins.json` (`docling-project/docling-layout-heron @ 8f39ad3c`,
//! `RTDetrV2ForObjectDetection`). Every stage is oracle-gated in `tests/layout_parity.rs`.
//!
//! Shape of the thing (all fixed by the pinned config):
//! * preproc: PIL-BILINEAR 640×640 + ×(1/255) — `vision::resize_rgb8_bilinear`, bit-exact;
//!   `do_normalize=false`, so the config's mean/std are DEAD — do not "fix" this.
//! * backbone: ResNet-50-D (`rt_detr_resnet`): deep 3-conv stem (relu) + maxpool; 4 bottleneck
//!   stages [3,4,6,3]; stride-2 lives in the 3×3 (`downsample_in_bottleneck=false`); stride-2
//!   shortcuts are `AvgPool(2,2,ceil) → 1×1` (ResNet-D); stage-1's shortcut is a plain 1×1.
//!   BN is frozen upstream (`replace_batch_norm`) — folded into the convs at load.
//! * hybrid encoder: AIFI (ONE post-norm layer, gelu-erf FFN, 2-D sincos pos on q/k) on the
//!   stride-32 map only (`encode_proj_layers=[2]`); CCFM FPN top-down + PAN bottom-up of
//!   CSPRep blocks (silu 1×1s + 3 RepVGG bottlenecks — 3×3+1×1 summed THEN silu; `conv3` is
//!   Identity at `hidden_expansion=1.0`), nearest-2× upsample, 3×3-s2 downsample.
//! * query selection: anchors `(grid+0.5)/size, wh=0.05·2^level` → logit space, invalid rows
//!   (outside eps=1e-2) forced to f32::MAX AND their memory rows zeroed; `enc_output` Linear+LN;
//!   topk-300 anchors by MAX CLASS LOGIT; decoder target = gathered `output_memory`
//!   (`learn_initial_query=false`); initial references = gathered coord LOGITS (pre-sigmoid).
//! * decoder: 6 post-norm layers — self-attn (q,k carry `query_pos_head(ref)` recomputed from the
//!   CURRENT reference each layer), v2 multi-scale deformable cross-attention (softmax over all
//!   `levels·points=12` jointly; offsets scaled by `n_points_scale · ref_wh · offset_scale=0.5`;
//!   the same cxcywh reference for every level; bilinear grid-sample, zeros padding,
//!   align_corners=false), relu FFN; iterative refinement in logit space
//!   (`sigmoid(bbox_embed(h) + inverse_sigmoid(ref))`).
//! * postproc: sigmoid scores, **topk-300 over the FLATTENED query×class grid** (one query may
//!   emit several labels), threshold 0.3, cxcywh→xyxy·[W,H,W,H], canonical docling label map
//!   (`label_offset=0` — the +1 shift belongs to the v1 architecture), clamp to page bounds.

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

use crate::conv2d::{Conv2d, fold_bn, grid_sample_bilinear, max_pool2d};
use crate::cpu_gemm::{PackedWeight, gemm_packed};
use crate::simd_math::erf_f32;
use crate::weights::LazySt;

/// One thresholded, labelled, page-space detection — the `MODEL_CONTRACT.md` §1 output row.
#[derive(Clone, Debug)]
pub struct Detection {
    pub label: String,
    pub confidence: f32,
    /// l, t, r, b in page-image pixel space, clamped to page bounds.
    pub bbox: [f32; 4],
}

/// Docling's canonical categories (`docling_ibm_models/layoutmodel/labels.py`) — the predictor
/// uses THIS map, not the checkpoint's `id2label`.
const CANONICAL_LABELS: [&str; 17] = [
    "Caption",
    "Footnote",
    "Formula",
    "List-item",
    "Page-footer",
    "Page-header",
    "Picture",
    "Section-header",
    "Table",
    "Text",
    "Title",
    "Document Index",
    "Code",
    "Checkbox-Selected",
    "Checkbox-Unselected",
    "Form",
    "Key-Value Region",
];

const D: usize = 256; // d_model
const HEADS: usize = 8;
const HDIM: usize = D / HEADS;
const LEVELS: usize = 3;
const POINTS: usize = 4;
const QUERIES: usize = 300;
const LABELS: usize = 17;
const DEC_LAYERS: usize = 6;
const LN_EPS: f32 = 1e-5;
const BN_EPS: f32 = 1e-5;
const THRESHOLD: f32 = 0.3;

// ── small pieces ────────────────────────────────────────────────────────────────────────────────

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

fn gelu_erf(x: f32) -> f32 {
    0.5 * x * (1.0 + erf_f32(x * std::f32::consts::FRAC_1_SQRT_2))
}

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

/// transformers' `inverse_sigmoid` (eps 1e-5), used by the iterative refinement.
fn inverse_sigmoid(x: f32) -> f32 {
    let x = x.clamp(0.0, 1.0);
    let x1 = x.max(1e-5);
    let x2 = (1.0 - x).max(1e-5);
    (x1 / x2).ln()
}

pub(crate) struct Lin {
    pub(crate) w: PackedWeight,
    pub(crate) 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 {
            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];
            }
        }
    }
}

#[derive(Clone, Copy, PartialEq)]
pub(crate) enum Act {
    None,
    Relu,
    Silu,
}

/// Conv + (folded) BN + activation — `RTDetrResNetConvLayer` / `RTDetrV2ConvNormLayer`.
pub(crate) struct ConvNorm {
    pub(crate) conv: Conv2d,
    pub(crate) act: Act,
}

impl ConvNorm {
    /// `conv_name` has `.weight` (no bias); `norm_name` carries eval BN to fold.
    fn load(
        st: &LazySt,
        conv_name: &str,
        norm_name: &str,
        stride: usize,
        act: Act,
    ) -> Result<Self> {
        let w = st.tensor_f32(&format!("{conv_name}.weight"))?;
        let shape = st.shape(&format!("{conv_name}.weight"))?.to_vec();
        let (oc, ic, kh, kw) = (shape[0], shape[1], shape[2], shape[3]);
        let gamma = st.tensor_f32(&format!("{norm_name}.weight"))?;
        let beta = st.tensor_f32(&format!("{norm_name}.bias"))?;
        let mean = st.tensor_f32(&format!("{norm_name}.running_mean"))?;
        let var = st.tensor_f32(&format!("{norm_name}.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),
            act,
        })
    }

    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (mut out, oh, ow) = self.conv.forward(x, h, w);
        match self.act {
            Act::None => {}
            Act::Relu => out.iter_mut().for_each(|v| *v = v.max(0.0)),
            Act::Silu => out.iter_mut().for_each(|v| *v = silu(*v)),
        }
        (out, oh, ow)
    }
}

/// NHWC 2×2/2 average pool, `ceil_mode=True` (ResNet-D shortcut). Even inputs only in this
/// model (160/80/40) — asserted, so ceil vs divisor subtleties never arise.
fn avg_pool2x2(x: &[f32], h: usize, w: usize, c: usize) -> (Vec<f32>, usize, usize) {
    assert!(h % 2 == 0 && w % 2 == 0, "avg_pool2x2: odd input {h}×{w}");
    let (oh, ow) = (h / 2, w / 2);
    let mut out = vec![0f32; oh * ow * c];
    for y in 0..oh {
        for xx in 0..ow {
            let dst = (y * ow + xx) * c;
            for (dy, dx) in [(0, 0), (0, 1), (1, 0), (1, 1)] {
                let src = ((2 * y + dy) * w + 2 * xx + dx) * c;
                for ch in 0..c {
                    out[dst + ch] += x[src + ch];
                }
            }
            for ch in 0..c {
                out[dst + ch] *= 0.25;
            }
        }
    }
    (out, oh, ow)
}

/// NHWC nearest-neighbour 2× upsample (`F.interpolate(scale_factor=2, mode="nearest")`).
fn upsample_nearest_2x(x: &[f32], h: usize, w: usize, c: usize) -> (Vec<f32>, usize, usize) {
    let (oh, ow) = (h * 2, w * 2);
    let mut out = vec![0f32; oh * ow * c];
    for y in 0..oh {
        for xx in 0..ow {
            let src = ((y / 2) * w + xx / 2) * c;
            let dst = (y * ow + xx) * c;
            out[dst..dst + c].copy_from_slice(&x[src..src + c]);
        }
    }
    (out, oh, ow)
}

/// Per-pixel channel concat of two NHWC maps with identical geometry.
fn concat_channels(a: &[f32], b: &[f32], hw: usize, ca: usize, cb: usize) -> Vec<f32> {
    let mut out = vec![0f32; hw * (ca + cb)];
    for p in 0..hw {
        out[p * (ca + cb)..p * (ca + cb) + ca].copy_from_slice(&a[p * ca..(p + 1) * ca]);
        out[p * (ca + cb) + ca..(p + 1) * (ca + cb)].copy_from_slice(&b[p * cb..(p + 1) * cb]);
    }
    out
}

// ── backbone ────────────────────────────────────────────────────────────────────────────────────

pub(crate) enum Shortcut {
    Identity,
    /// plain 1×1 (+BN) — stage 1's channel-only projection
    Conv(ConvNorm),
    /// ResNet-D: AvgPool(2,2,ceil) → 1×1 stride-1 (+BN)
    PoolConv(ConvNorm),
}

pub(crate) struct Bottleneck {
    pub(crate) layer: [ConvNorm; 3], // 1×1 reduce (relu), 3×3 [stride] (relu), 1×1 expand (none)
    pub(crate) shortcut: Shortcut,
}

impl Bottleneck {
    fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (a, h1, w1) = self.layer[0].forward(x, h, w);
        let (a, h2, w2) = self.layer[1].forward(&a, h1, w1);
        let (mut a, h3, w3) = self.layer[2].forward(&a, h2, w2);
        // residual add + relu, reading the shortcut source in place — an Identity shortcut
        // copying its whole feature map was measurable backbone wall time (profiled)
        let owned: Option<Vec<f32>> = match &self.shortcut {
            Shortcut::Identity => None,
            Shortcut::Conv(sc) => Some(sc.forward(x, h, w).0),
            Shortcut::PoolConv(sc) => {
                let (p, ph, pw) = avg_pool2x2(x, h, w, x.len() / (h * w));
                Some(sc.forward(&p, ph, pw).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, h3, w3)
    }
}

pub(crate) struct Backbone {
    pub(crate) stem: [ConvNorm; 3],
    pub(crate) stages: Vec<Vec<Bottleneck>>,
}

impl Backbone {
    fn load(st: &LazySt) -> Result<Self> {
        let stem_name = |i: usize| format!("model.backbone.model.embedder.embedder.{i}");
        let stem = [
            ConvNorm::load(
                st,
                &format!("{}.convolution", stem_name(0)),
                &format!("{}.normalization", stem_name(0)),
                2,
                Act::Relu,
            )?,
            ConvNorm::load(
                st,
                &format!("{}.convolution", stem_name(1)),
                &format!("{}.normalization", stem_name(1)),
                1,
                Act::Relu,
            )?,
            ConvNorm::load(
                st,
                &format!("{}.convolution", stem_name(2)),
                &format!("{}.normalization", stem_name(2)),
                1,
                Act::Relu,
            )?,
        ];

        let depths = [3usize, 4, 6, 3];
        let mut stages = Vec::with_capacity(4);
        for (s, &depth) in depths.iter().enumerate() {
            let mut blocks = Vec::with_capacity(depth);
            for l in 0..depth {
                let base = format!("model.backbone.model.encoder.stages.{s}.layers.{l}");
                let conv = |i: usize, stride: usize, act: Act| {
                    ConvNorm::load(
                        st,
                        &format!("{base}.layer.{i}.convolution"),
                        &format!("{base}.layer.{i}.normalization"),
                        stride,
                        act,
                    )
                };
                // stride-2 lives in the 3×3 (downsample_in_bottleneck=false); stage 0 is stride 1
                let stride = if l == 0 && s > 0 { 2 } else { 1 };
                let layer = [
                    conv(0, 1, Act::Relu)?,
                    conv(1, stride, Act::Relu)?,
                    conv(2, 1, Act::None)?,
                ];
                let shortcut = if l == 0 {
                    if s == 0 {
                        // stage 1: channel projection only, plain 1×1 (checkpoint: `shortcut.convolution`)
                        Shortcut::Conv(ConvNorm::load(
                            st,
                            &format!("{base}.shortcut.convolution"),
                            &format!("{base}.shortcut.normalization"),
                            1,
                            Act::None,
                        )?)
                    } else {
                        // ResNet-D Sequential[avgpool, shortcut] (checkpoint: `shortcut.1.convolution`)
                        Shortcut::PoolConv(ConvNorm::load(
                            st,
                            &format!("{base}.shortcut.1.convolution"),
                            &format!("{base}.shortcut.1.normalization"),
                            1,
                            Act::None,
                        )?)
                    }
                } else {
                    Shortcut::Identity
                };
                blocks.push(Bottleneck { layer, shortcut });
            }
            stages.push(blocks);
        }
        Ok(Self { stem, stages })
    }

    /// 640×640×3 NHWC → the three feature maps the encoder consumes (strides 8/16/32).
    fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<(Vec<f32>, usize, usize)> {
        let (x, h, w) = self.stem[0].forward(x, h, w);
        let (x, h, w) = self.stem[1].forward(&x, h, w);
        let (x, h, w) = self.stem[2].forward(&x, h, w);
        let (mut x, mut h, mut w) = max_pool2d(&x, h, w, 64, 3, 2, 1);
        let mut out = Vec::with_capacity(3);
        for (s, blocks) in self.stages.iter().enumerate() {
            for b in blocks {
                let (nx, nh, nw) = b.forward(&x, h, w);
                x = nx;
                h = nh;
                w = nw;
            }
            if s >= 1 {
                out.push((x.clone(), h, w)); // stages 2,3,4 → strides 8,16,32
            }
        }
        out
    }
}

// ── hybrid encoder (AIFI + CCFM) ────────────────────────────────────────────────────────────────

/// One post-norm transformer encoder layer (AIFI): q,k carry the sincos pos embedding.
pub(crate) struct EncLayer {
    q: Lin,
    k: Lin,
    v: Lin,
    out: Lin,
    ln1: LayerNorm,
    fc1: Lin,
    fc2: Lin,
    ln2: LayerNorm,
}

impl EncLayer {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        Ok(Self {
            q: Lin::load(st, &format!("{base}.self_attn.q_proj"))?,
            k: Lin::load(st, &format!("{base}.self_attn.k_proj"))?,
            v: Lin::load(st, &format!("{base}.self_attn.v_proj"))?,
            out: Lin::load(st, &format!("{base}.self_attn.out_proj"))?,
            ln1: LayerNorm::load(st, &format!("{base}.self_attn_layer_norm"))?,
            fc1: Lin::load(st, &format!("{base}.fc1"))?,
            fc2: Lin::load(st, &format!("{base}.fc2"))?,
            ln2: LayerNorm::load(st, &format!("{base}.final_layer_norm"))?,
        })
    }

    pub(crate) fn forward(&self, src: &[f32], pos: &[f32], n: usize) -> Vec<f32> {
        // q = k = src + pos; v = src (DETR convention)
        let withpos: Vec<f32> = src.iter().zip(pos).map(|(a, b)| a + b).collect();
        let q = self.q.forward(&withpos, n);
        let k = self.k.forward(&withpos, n);
        let v = self.v.forward(src, n);
        let attn = mha(&q, &k, &v, n, n);
        let attn = self.out.forward(&attn, n);
        let mut h: Vec<f32> = src.iter().zip(&attn).map(|(a, b)| a + b).collect();
        self.ln1.forward_inplace(&mut h);
        let mut f = self.fc1.forward(&h, n);
        f.iter_mut().for_each(|x| *x = gelu_erf(*x)); // encoder_activation_function = gelu
        let f = self.fc2.forward(&f, n);
        let mut o: Vec<f32> = h.iter().zip(&f).map(|(a, b)| a + b).collect();
        self.ln2.forward_inplace(&mut o);
        o
    }
}

/// Plain multi-head attention `[nq, D] × [nk, D]` → `[nq, D]` (8 heads × 32, 1/√32 scaling).
///
/// Same perf shape as tableformer's `mha512` (the decode-phase trace put the decoder's
/// `pos+self` at ~50 ms — this loop, scalar and serial): rayon over QUERY rows (per-row
/// arithmetic order unchanged → bit-identical to the serial loop) + `simd::dot` for the QKᵀ
/// reductions (NEON reassociation — the 77/77-detections gate is the arbiter).
fn mha(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];
    out.par_chunks_mut(D).enumerate().for_each(|(i, orow)| {
        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;
                }
            }
        }
    });
    out
}

/// 2-D sincos position embedding, exactly `build_2d_sincos_position_embedding` (temperature 1e4).
pub(crate) fn sincos_pos_embed(w: usize, h: usize) -> Vec<f32> {
    let pos_dim = D / 4; // 64
    let omega: Vec<f32> = (0..pos_dim)
        .map(|i| 1.0 / 10000f32.powf(i as f32 / pos_dim as f32))
        .collect();
    let mut out = vec![0f32; w * h * D];
    // torch builds [w*h] rows with meshgrid(w, h, indexing="ij"): row t has grid_w = t/h, grid_h = t%h.
    for t in 0..w * h {
        let gw = (t / h) as f32;
        let gh = (t % h) as f32;
        let row = &mut out[t * D..(t + 1) * D];
        for (i, &om) in omega.iter().enumerate() {
            row[i] = (gw * om).sin();
            row[pos_dim + i] = (gw * om).cos();
            row[2 * pos_dim + i] = (gh * om).sin();
            row[3 * pos_dim + i] = (gh * om).cos();
        }
    }
    out
}

pub(crate) struct RepVgg {
    pub(crate) conv1: ConvNorm, // 3×3, no act
    pub(crate) conv2: ConvNorm, // 1×1, no act
}

impl RepVgg {
    fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
        let (a, ..) = self.conv1.forward(x, h, w);
        let (b, ..) = self.conv2.forward(x, h, w);
        a.iter()
            .zip(&b)
            .map(|(p, q)| silu(p + q)) // activation AFTER the sum
            .collect()
    }
}

pub(crate) struct Csp {
    pub(crate) conv1: ConvNorm, // 1×1 silu, in=512
    pub(crate) conv2: ConvNorm, // 1×1 silu, in=512
    pub(crate) bottlenecks: Vec<RepVgg>,
    // conv3 = Identity (hidden_expansion = 1.0)
}

impl Csp {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        let cn = |name: &str, stride: usize, act: Act| {
            ConvNorm::load(
                st,
                &format!("{base}.{name}.conv"),
                &format!("{base}.{name}.norm"),
                stride,
                act,
            )
        };
        let mut bottlenecks = Vec::with_capacity(3);
        for i in 0..3 {
            bottlenecks.push(RepVgg {
                conv1: cn(&format!("bottlenecks.{i}.conv1"), 1, Act::None)?,
                conv2: cn(&format!("bottlenecks.{i}.conv2"), 1, Act::None)?,
            });
        }
        Ok(Self {
            conv1: cn("conv1", 1, Act::Silu)?,
            conv2: cn("conv2", 1, Act::Silu)?,
            bottlenecks,
        })
    }

    fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
        let (mut a, ..) = self.conv1.forward(x, h, w);
        for b in &self.bottlenecks {
            a = b.forward(&a, h, w);
        }
        let (c, ..) = self.conv2.forward(x, h, w);
        a.iter_mut().zip(&c).for_each(|(p, q)| *p += q);
        a // conv3 is Identity
    }
}

pub(crate) struct HybridEncoder {
    pub(crate) input_proj: Vec<(Conv2d, ())>, // 1×1 + BN folded, NO activation
    pub(crate) aifi: EncLayer,
    pub(crate) lateral: Vec<ConvNorm>,
    pub(crate) fpn: Vec<Csp>,
    pub(crate) downsample: Vec<ConvNorm>,
    pub(crate) pan: Vec<Csp>,
}

impl HybridEncoder {
    fn load(st: &LazySt) -> Result<Self> {
        let mut input_proj = Vec::with_capacity(3);
        for i in 0..3 {
            let cn = ConvNorm::load(
                st,
                &format!("model.encoder_input_proj.{i}.0"),
                &format!("model.encoder_input_proj.{i}.1"),
                1,
                Act::None,
            )?;
            input_proj.push((cn.conv, ()));
        }
        let mut lateral = Vec::new();
        let mut fpn = Vec::new();
        let mut downsample = Vec::new();
        let mut pan = Vec::new();
        for i in 0..2 {
            lateral.push(ConvNorm::load(
                st,
                &format!("model.encoder.lateral_convs.{i}.conv"),
                &format!("model.encoder.lateral_convs.{i}.norm"),
                1,
                Act::Silu,
            )?);
            fpn.push(Csp::load(st, &format!("model.encoder.fpn_blocks.{i}"))?);
            downsample.push(ConvNorm::load(
                st,
                &format!("model.encoder.downsample_convs.{i}.conv"),
                &format!("model.encoder.downsample_convs.{i}.norm"),
                2,
                Act::Silu,
            )?);
            pan.push(Csp::load(st, &format!("model.encoder.pan_blocks.{i}"))?);
        }
        Ok(Self {
            input_proj,
            aifi: EncLayer::load(st, "model.encoder.encoder.0.layers.0")?,
            lateral,
            fpn,
            downsample,
            pan,
        })
    }

    /// 3 backbone maps → 3 fused maps (strides 8/16/32), all 256-channel NHWC.
    fn forward(&self, feats: &[(Vec<f32>, usize, usize)]) -> Vec<(Vec<f32>, usize, usize)> {
        // project to 256
        let mut maps: Vec<(Vec<f32>, usize, usize)> = feats
            .iter()
            .zip(&self.input_proj)
            .map(|((f, h, w), (proj, ()))| {
                let (m, oh, ow) = proj.forward(f, *h, *w);
                (m, oh, ow)
            })
            .collect();

        // AIFI on the stride-32 map (encode_proj_layers = [2]); NHWC IS token-major [h*w, 256]
        {
            let (m, h, w) = &maps[2];
            let pos = sincos_pos_embed(*w, *h);
            let o = self.aifi.forward(m, &pos, h * w);
            maps[2] = (o, *h, *w);
        }

        // top-down FPN: fpn_maps[-1] is REPLACED by its lateral conv each iteration
        let mut fpn_maps: Vec<(Vec<f32>, usize, usize)> = vec![maps[2].clone()];
        for idx in 0..2 {
            let (bf, bh, bw) = &maps[1 - idx]; // stride 16, then stride 8
            let (top, th, tw) = fpn_maps.last().unwrap().clone();
            let (lat, ..) = self.lateral[idx].forward(&top, th, tw);
            *fpn_maps.last_mut().unwrap() = (lat.clone(), th, tw);
            let (up, uh, uw) = upsample_nearest_2x(&lat, th, tw, D);
            debug_assert_eq!((uh, uw), (*bh, *bw));
            let fused = concat_channels(&up, bf, uh * uw, D, D);
            let o = self.fpn[idx].forward(&fused, uh, uw);
            fpn_maps.push((o, uh, uw));
        }
        fpn_maps.reverse(); // [s8, lat(s16-fused), lat(s32)]

        // bottom-up PAN
        let mut pan_maps: Vec<(Vec<f32>, usize, usize)> = vec![fpn_maps[0].clone()];
        for idx in 0..2 {
            let (top, th, tw) = pan_maps.last().unwrap().clone();
            let (down, dh, dw) = self.downsample[idx].forward(&top, th, tw);
            let (ff, fh, fw) = &fpn_maps[idx + 1];
            debug_assert_eq!((dh, dw), (*fh, *fw));
            let fused = concat_channels(&down, ff, dh * dw, D, D);
            let o = self.pan[idx].forward(&fused, dh, dw);
            pan_maps.push((o, dh, dw));
        }
        pan_maps
    }
}

// ── decoder ─────────────────────────────────────────────────────────────────────────────────────

struct Mlp {
    layers: Vec<Lin>,
}

impl Mlp {
    fn load(st: &LazySt, base: &str, n_layers: usize) -> Result<Self> {
        let mut layers = Vec::with_capacity(n_layers);
        for i in 0..n_layers {
            layers.push(Lin::load(st, &format!("{base}.layers.{i}"))?);
        }
        Ok(Self { layers })
    }

    /// relu between layers, none after the last.
    fn forward(&self, x: &[f32], m: usize) -> Vec<f32> {
        let mut h = x.to_vec();
        for (i, l) in self.layers.iter().enumerate() {
            h = l.forward(&h, m);
            if i + 1 < self.layers.len() {
                h.iter_mut().for_each(|v| *v = v.max(0.0));
            }
        }
        h
    }
}

pub(crate) struct DecLayer {
    q: Lin,
    k: Lin,
    v: Lin,
    out: Lin,
    ln1: LayerNorm,
    sampling_offsets: Lin,
    attention_weights: Lin,
    pub(crate) value_proj: Lin,
    output_proj: Lin,
    n_points_scale: Vec<f32>,
    ln2: LayerNorm,
    fc1: Lin,
    fc2: Lin,
    ln3: LayerNorm,
}

impl DecLayer {
    fn load(st: &LazySt, base: &str) -> Result<Self> {
        Ok(Self {
            q: Lin::load(st, &format!("{base}.self_attn.q_proj"))?,
            k: Lin::load(st, &format!("{base}.self_attn.k_proj"))?,
            v: Lin::load(st, &format!("{base}.self_attn.v_proj"))?,
            out: Lin::load(st, &format!("{base}.self_attn.out_proj"))?,
            ln1: LayerNorm::load(st, &format!("{base}.self_attn_layer_norm"))?,
            sampling_offsets: Lin::load(st, &format!("{base}.encoder_attn.sampling_offsets"))?,
            attention_weights: Lin::load(st, &format!("{base}.encoder_attn.attention_weights"))?,
            value_proj: Lin::load(st, &format!("{base}.encoder_attn.value_proj"))?,
            output_proj: Lin::load(st, &format!("{base}.encoder_attn.output_proj"))?,
            n_points_scale: st.tensor_f32(&format!("{base}.encoder_attn.n_points_scale"))?,
            ln2: LayerNorm::load(st, &format!("{base}.encoder_attn_layer_norm"))?,
            fc1: Lin::load(st, &format!("{base}.fc1"))?,
            fc2: Lin::load(st, &format!("{base}.fc2"))?,
            ln3: LayerNorm::load(st, &format!("{base}.final_layer_norm"))?,
        })
    }
}

pub(crate) struct Decoder {
    pub(crate) layers: Vec<DecLayer>,
    query_pos_head: Mlp,
    bbox_embed: Vec<Mlp>,
    class_embed: Vec<Lin>,
}

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

/// Everything `forward_stages` exposes, mirroring the oracle's stage keys one-to-one so the
/// parity test can gate each piece in isolation.
pub struct Stages {
    /// 3 backbone maps (NHWC, h, w) — strides 8/16/32.
    pub backbone: Vec<(Vec<f32>, usize, usize)>,
    /// 3 fused encoder maps.
    pub encoder: Vec<(Vec<f32>, usize, usize)>,
    /// `[anchors_total, LABELS]` class logits over the flattened memory.
    pub enc_outputs_class: Vec<f32>,
    /// `[anchors_total, 4]` coord logits (bbox head + anchors).
    pub enc_outputs_coord_logits: Vec<f32>,
    /// topk-300 anchor indices, torch `topk` order (score desc, ties by lower index).
    pub topk_idx: Vec<usize>,
    /// `[300, 4]` initial reference logits (pre-sigmoid).
    pub init_reference_logits: Vec<f32>,
    /// per decoder layer: hidden `[300, D]`.
    pub dec_hidden: Vec<Vec<f32>>,
    /// per decoder layer: POST-refinement reference points, sigmoid space `[300, 4]`.
    pub dec_refs: Vec<Vec<f32>>,
    /// per decoder layer: class logits `[300, LABELS]`.
    pub dec_logits: Vec<Vec<f32>>,
}

pub struct RtDetr {
    pub(crate) backbone: Backbone,
    pub(crate) encoder: HybridEncoder,
    pub(crate) dec_input_proj: Vec<Conv2d>,
    enc_output: Lin,
    enc_output_ln: LayerNorm,
    enc_score_head: Lin,
    enc_bbox_head: Mlp,
    pub(crate) decoder: Decoder,
}

impl RtDetr {
    pub fn load(dir: &Path) -> Result<Self> {
        let st = LazySt::open(dir).context("open layout checkpoint")?;
        let backbone = Backbone::load(&st)?;
        let encoder = HybridEncoder::load(&st)?;
        let mut dec_input_proj = Vec::with_capacity(3);
        for i in 0..3 {
            let cn = ConvNorm::load(
                &st,
                &format!("model.decoder_input_proj.{i}.0"),
                &format!("model.decoder_input_proj.{i}.1"),
                1,
                Act::None,
            )?;
            dec_input_proj.push(cn.conv);
        }
        let mut layers = Vec::with_capacity(DEC_LAYERS);
        let mut bbox_embed = Vec::with_capacity(DEC_LAYERS);
        let mut class_embed = Vec::with_capacity(DEC_LAYERS);
        for i in 0..DEC_LAYERS {
            layers.push(DecLayer::load(&st, &format!("model.decoder.layers.{i}"))?);
            bbox_embed.push(Mlp::load(&st, &format!("model.decoder.bbox_embed.{i}"), 3)?);
            class_embed.push(Lin::load(&st, &format!("model.decoder.class_embed.{i}"))?);
        }
        Ok(Self {
            backbone,
            encoder,
            dec_input_proj,
            enc_output: Lin::load(&st, "model.enc_output.0")?,
            enc_output_ln: LayerNorm::load(&st, "model.enc_output.1")?,
            enc_score_head: Lin::load(&st, "model.enc_score_head")?,
            enc_bbox_head: Mlp::load(&st, "model.enc_bbox_head", 3)?,
            decoder: Decoder {
                layers,
                query_pos_head: Mlp::load(&st, "model.decoder.query_pos_head", 2)?,
                bbox_embed,
                class_embed,
            },
        })
    }

    /// Full forward with every stage exposed. `pixel_values` is NHWC 640×640×3, already
    /// preprocessed (bilinear-resized u8 × 1/255).
    ///
    /// `OSFKB_RTDETR_TRACE=1` prints per-stage wall times (house profiling convention — every
    /// perf lever here was picked from this trace, not guessed).
    pub fn forward_stages(&self, pixel_values: &[f32], h: usize, w: usize) -> Stages {
        let trace = std::env::var("OSFKB_RTDETR_TRACE").is_ok();
        let mut t = std::time::Instant::now();
        let mut lap = |name: &str| {
            if trace {
                eprintln!("rtdetr {name}: {:.1} ms", t.elapsed().as_secs_f64() * 1e3);
            }
            t = std::time::Instant::now();
        };
        let backbone = self.backbone.forward(pixel_values, h, w);
        lap("backbone");
        let encoder = self.encoder.forward(&backbone);
        lap("hybrid-encoder");
        let s = self.stages_from_maps(backbone, encoder);
        lap("qs+decoder");
        return s;
    }

    /// Query selection + decoder from already-computed encoder maps — the shared tail the GPU
    /// backbone/encoder path (`rtdetr_gpu`) hands its maps to.
    pub(crate) fn stages_from_maps(
        &self,
        backbone: Vec<(Vec<f32>, usize, usize)>,
        encoder: Vec<(Vec<f32>, usize, usize)>,
    ) -> Stages {
        // decoder input projection + flatten (level order s8, s16, s32)
        let mut memory: Vec<f32> = Vec::new();
        let mut shapes: Vec<(usize, usize)> = Vec::new();
        for (i, (m, mh, mw)) in encoder.iter().enumerate() {
            let (p, ..) = self.dec_input_proj[i].forward(m, *mh, *mw);
            memory.extend_from_slice(&p);
            shapes.push((*mh, *mw));
        }
        self.stages_from_projected(backbone, encoder, memory, shapes, None)
    }

    /// The tail from an already-projected flattened memory — the GPU path computes
    /// `dec_input_proj` and (optionally) every decoder layer's `value_proj` on-device and hands
    /// the results here, so both arms share ALL remaining logic.
    pub(crate) fn stages_from_projected(
        &self,
        backbone: Vec<(Vec<f32>, usize, usize)>,
        encoder: Vec<(Vec<f32>, usize, usize)>,
        memory: Vec<f32>,
        shapes: Vec<(usize, usize)>,
        values: Option<Vec<Vec<f32>>>,
    ) -> Stages {
        let total: usize = shapes.iter().map(|(a, b)| a * b).sum();

        // anchors in logit space + validity
        let mut anchors = vec![0f32; total * 4];
        let mut valid = vec![true; total];
        {
            let mut off = 0usize;
            for (lvl, &(gh, gw)) in shapes.iter().enumerate() {
                let wh = 0.05f32 * 2f32.powi(lvl as i32);
                for y in 0..gh {
                    for x in 0..gw {
                        let i = off + y * gw + x;
                        let a = [
                            (x as f32 + 0.5) / gw as f32,
                            (y as f32 + 0.5) / gh as f32,
                            wh,
                            wh,
                        ];
                        let ok = a.iter().all(|&v| v > 1e-2 && v < 1.0 - 1e-2);
                        valid[i] = ok;
                        for (j, &v) in a.iter().enumerate() {
                            anchors[i * 4 + j] = if ok { (v / (1.0 - v)).ln() } else { f32::MAX };
                        }
                    }
                }
                off += gh * gw;
            }
        }

        // memory rows for invalid anchors are ZEROED before enc_output (torch: valid_mask * memory)
        let mut masked = memory.clone();
        for (i, &ok) in valid.iter().enumerate() {
            if !ok {
                masked[i * D..(i + 1) * D].fill(0.0);
            }
        }
        let mut output_memory = self.enc_output.forward(&masked, total);
        self.enc_output_ln.forward_inplace(&mut output_memory);

        let enc_outputs_class = self.enc_score_head.forward(&output_memory, total);
        let mut enc_outputs_coord_logits = self.enc_bbox_head.forward(&output_memory, total);
        for (i, v) in enc_outputs_coord_logits.iter_mut().enumerate() {
            // f32::MAX + finite overflows to inf; torch keeps finfo.max (adds then stays finite
            // only by accident of magnitudes) — clamp to the torch value for parity.
            *v += anchors[i];
            if !v.is_finite() {
                *v = f32::MAX;
            }
        }

        // topk-300 anchors by max class logit; torch topk = score desc, ties by lower index
        let mut best = vec![f32::NEG_INFINITY; total];
        for i in 0..total {
            for c in 0..LABELS {
                best[i] = best[i].max(enc_outputs_class[i * LABELS + c]);
            }
        }
        let mut order: Vec<usize> = (0..total).collect();
        order.sort_by(|&a, &b| best[b].partial_cmp(&best[a]).unwrap().then(a.cmp(&b)));
        let topk_idx: Vec<usize> = order[..QUERIES].to_vec();

        let mut init_reference_logits = vec![0f32; QUERIES * 4];
        let mut target = vec![0f32; QUERIES * D];
        for (q, &i) in topk_idx.iter().enumerate() {
            init_reference_logits[q * 4..q * 4 + 4]
                .copy_from_slice(&enc_outputs_coord_logits[i * 4..i * 4 + 4]);
            target[q * D..(q + 1) * D].copy_from_slice(&output_memory[i * D..(i + 1) * D]);
        }

        // per-layer deformable value gathers work on the UNMASKED flattened memory
        let (dec_hidden, dec_refs, dec_logits) = self.decode(
            &target,
            &init_reference_logits,
            &memory,
            &shapes,
            values.as_deref(),
        );

        Stages {
            backbone,
            encoder,
            enc_outputs_class,
            enc_outputs_coord_logits,
            topk_idx,
            init_reference_logits,
            dec_hidden,
            dec_refs,
            dec_logits,
        }
    }

    #[allow(clippy::type_complexity)]
    fn decode(
        &self,
        target: &[f32],
        init_ref_logits: &[f32],
        memory: &[f32],
        shapes: &[(usize, usize)],
        precomputed_values: Option<&[Vec<f32>]>,
    ) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<Vec<f32>>) {
        let total: usize = shapes.iter().map(|(a, b)| a * b).sum();
        let mut hidden = target.to_vec();
        let mut refs: Vec<f32> = init_ref_logits.iter().map(|&v| sigmoid(v)).collect();
        let mut all_hidden = Vec::with_capacity(DEC_LAYERS);
        let mut all_refs = Vec::with_capacity(DEC_LAYERS);
        let mut all_logits = Vec::with_capacity(DEC_LAYERS);

        // OSFKB_RTDETR_TRACE=2 additionally splits the decoder internals (per-phase totals over
        // all 6 layers) — the margin-widening iterations pick their targets from this, not vibes.
        let trace2 = std::env::var("OSFKB_RTDETR_TRACE").is_ok_and(|v| v == "2");
        let mut phase = [0f64; 5]; // pos+self / value / sampling / ffn+ln / heads
        macro_rules! ph {
            ($i:expr, $t:expr) => {
                if trace2 {
                    phase[$i] += $t.elapsed().as_secs_f64() * 1e3;
                }
            };
        }

        for (li, layer) in self.decoder.layers.iter().enumerate() {
            let t0 = std::time::Instant::now();
            // query pos from CURRENT refs (sigmoid space)
            let pos = self.decoder.query_pos_head.forward(&refs, QUERIES);

            // self-attention (post-norm)
            let withpos: Vec<f32> = hidden.iter().zip(&pos).map(|(a, b)| a + b).collect();
            let q = layer.q.forward(&withpos, QUERIES);
            let k = layer.k.forward(&withpos, QUERIES);
            let v = layer.v.forward(&hidden, QUERIES);
            let sa = layer
                .out
                .forward(&mha(&q, &k, &v, QUERIES, QUERIES), QUERIES);
            let mut h1: Vec<f32> = hidden.iter().zip(&sa).map(|(a, b)| a + b).collect();
            layer.ln1.forward_inplace(&mut h1);
            ph!(0, t0);
            let t0 = std::time::Instant::now();

            // deformable cross-attention: pos is added to the QUERY input here too
            let hq: Vec<f32> = h1.iter().zip(&pos).map(|(a, b)| a + b).collect();
            let offsets = layer.sampling_offsets.forward(&hq, QUERIES); // [q, 8·12·2]
            let mut weights = layer.attention_weights.forward(&hq, QUERIES); // [q, 8·12]
            // softmax over the joint levels×points axis, per head
            for qi in 0..QUERIES {
                for hd in 0..HEADS {
                    let s = &mut weights[qi * HEADS * LEVELS * POINTS + hd * LEVELS * POINTS..]
                        [..LEVELS * POINTS];
                    let m = s.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
                    let mut den = 0f32;
                    for v in s.iter_mut() {
                        *v = (*v - m).exp();
                        den += *v;
                    }
                    s.iter_mut().for_each(|v| *v /= den);
                }
            }
            // [total, D] — GPU-precomputed when the wgpu arm ran value_proj on-device
            let value: std::borrow::Cow<[f32]> = match precomputed_values {
                Some(vs) => std::borrow::Cow::Borrowed(&vs[li]),
                None => std::borrow::Cow::Owned(layer.value_proj.forward(memory, total)),
            };
            ph!(1, t0);
            let t0 = std::time::Instant::now();

            // per (level, head): contiguous [H·W, HDIM] slice, then bilinear reads
            let mut cross = vec![0f32; QUERIES * D];
            let mut lvl_off = 0usize;
            for (lvl, &(gh, gw)) in shapes.iter().enumerate() {
                let hw = gh * gw;
                for hd in 0..HEADS {
                    let mut slice = vec![0f32; hw * HDIM];
                    for p in 0..hw {
                        let src = (lvl_off + p) * D + hd * HDIM;
                        slice[p * HDIM..(p + 1) * HDIM].copy_from_slice(&value[src..src + HDIM]);
                    }
                    let mut grid = Vec::with_capacity(QUERIES * POINTS);
                    for qi in 0..QUERIES {
                        let r = &refs[qi * 4..qi * 4 + 4];
                        for pt in 0..POINTS {
                            let k = lvl * POINTS + pt;
                            let o =
                                qi * HEADS * LEVELS * POINTS * 2 + (hd * LEVELS * POINTS + k) * 2;
                            let scale = layer.n_points_scale[k] * 0.5; // n_points_scale · offset_scale
                            let lx = r[0] + offsets[o] * scale * r[2];
                            let ly = r[1] + offsets[o + 1] * scale * r[3];
                            grid.push((2.0 * lx - 1.0, 2.0 * ly - 1.0));
                        }
                    }
                    let sampled = grid_sample_bilinear(&slice, gh, gw, HDIM, &grid);
                    for qi in 0..QUERIES {
                        for pt in 0..POINTS {
                            let k = lvl * POINTS + pt;
                            let wgt =
                                weights[qi * HEADS * LEVELS * POINTS + hd * LEVELS * POINTS + k];
                            let s = &sampled[(qi * POINTS + pt) * HDIM..][..HDIM];
                            let o = &mut cross[qi * D + hd * HDIM..][..HDIM];
                            for (ov, sv) in o.iter_mut().zip(s) {
                                *ov += wgt * sv;
                            }
                        }
                    }
                }
                lvl_off += hw;
            }
            let ca = layer.output_proj.forward(&cross, QUERIES);
            let mut h2: Vec<f32> = h1.iter().zip(&ca).map(|(a, b)| a + b).collect();
            layer.ln2.forward_inplace(&mut h2);

            ph!(2, t0);
            let t0 = std::time::Instant::now();
            // FFN (relu)
            let mut f = layer.fc1.forward(&h2, QUERIES);
            f.iter_mut().for_each(|x| *x = x.max(0.0));
            let f = layer.fc2.forward(&f, QUERIES);
            let mut h3: Vec<f32> = h2.iter().zip(&f).map(|(a, b)| a + b).collect();
            layer.ln3.forward_inplace(&mut h3);
            hidden = h3;

            // iterative refinement in logit space
            let corners = self.decoder.bbox_embed[li].forward(&hidden, QUERIES);
            let mut new_refs = vec![0f32; QUERIES * 4];
            for i in 0..QUERIES * 4 {
                new_refs[i] = sigmoid(corners[i] + inverse_sigmoid(refs[i]));
            }
            refs = new_refs;

            ph!(3, t0);
            let t0 = std::time::Instant::now();
            all_logits.push(self.decoder.class_embed[li].forward(&hidden, QUERIES));
            all_hidden.push(hidden.clone());
            all_refs.push(refs.clone());
            ph!(4, t0);
        }
        if trace2 {
            eprintln!(
                "rtdetr decode phases (ms over 6 layers): pos+self={:.1} value={:.1} sample+cross={:.1} ffn={:.1} heads+refine={:.1}",
                phase[0], phase[1], phase[2], phase[3], phase[4]
            );
        }
        (all_hidden, all_refs, all_logits)
    }

    /// The `MODEL_CONTRACT.md` §1 surface: page RGB8 → thresholded detections in page space.
    pub fn detect(&self, rgb: &[u8], w: usize, h: usize) -> Vec<Detection> {
        let resized = crate::vision::resize_rgb8_bilinear(rgb, w, h, 640, 640);
        let pixels: Vec<f32> = resized.iter().map(|&b| b as f32 / 255.0).collect();
        let stages = self.forward_stages(&pixels, 640, 640);
        postprocess(
            stages.dec_logits.last().unwrap(),
            stages.dec_refs.last().unwrap(),
            w as f32,
            h as f32,
        )
    }
}

/// `post_process_object_detection(threshold=0.3)` + the predictor's label map and clamping.
pub fn postprocess(
    logits: &[f32],
    boxes_cxcywh: &[f32],
    page_w: f32,
    page_h: f32,
) -> Vec<Detection> {
    // sigmoid scores over the FLATTENED [query × class] grid, topk-300
    let mut flat: Vec<(f32, usize)> = (0..QUERIES * LABELS)
        .map(|i| (sigmoid(logits[i]), i))
        .collect();
    flat.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap().then(a.1.cmp(&b.1)));
    let mut out = Vec::new();
    for &(score, idx) in flat.iter().take(QUERIES) {
        if score <= THRESHOLD {
            continue;
        }
        let (q, label) = (idx / LABELS, idx % LABELS);
        let b = &boxes_cxcywh[q * 4..q * 4 + 4];
        let (cx, cy, bw, bh) = (b[0] * page_w, b[1] * page_h, b[2] * page_w, b[3] * page_h);
        out.push(Detection {
            label: CANONICAL_LABELS[label].to_string(),
            confidence: score,
            bbox: [
                (cx - 0.5 * bw).clamp(0.0, page_w),
                (cy - 0.5 * bh).clamp(0.0, page_h),
                (cx + 0.5 * bw).clamp(0.0, page_w),
                (cy + 0.5 * bh).clamp(0.0, page_h),
            ],
        });
    }
    out
}