inferencelayer 0.2.3

//! DeepSeek-OCR **DeepEncoder** — the optical-compression vision encoder, as a CPU reference.
//!
//! DeepEncoder turns a 1024×1024 page into **256 vision tokens** (16× fewer than a dense-text
//! transcription would cost), by running the expensive dense attention only on the *compressed*
//! grid. The pipeline, with exact dims from `deepseek-ai/DeepSeek-OCR/config.json`:
//!
//! ```text
//!   1024²×3 ─PatchEmbed(conv16,s16)→ 64×64×768 ─SAM(12 blk, win14, global@[2,5,8,11])→ 64×64×768
//!          ─neck(768→256)→ 64×64×256 ─compress(2×conv k3 s2 p1: 256→512→1024)→ 16×16×1024
//!          ─CLIP-L(24 blk, global)→ 256×1024 ─proj(concat 2048→1280)→ 256 tokens
//! ```
//!
//! **Build discipline** (mirrors [`crate::vision_gpu`]): this module is the numeric ORACLE. It is
//! written for clarity, not speed — plain loops, f32 throughout — so the GPU kernels in
//! `deepencoder_gpu` (step 2) can be gated against it, and this module in turn against the DeepSeek
//! torch reference. Nothing here is on the serving hot path.
//!
//! **Status: STEP 1** — geometry, preprocessing, patch-embed, and the SAM block (incl. the
//! decomposed relative-position attention, DeepEncoder's one non-standard kernel) are implemented
//! and shape-tested. The neck/compressor/CLIP/projector stages carry exact specs and are wired as
//! typed stages; each is filled + torch-gated in a following step. See the module test module for
//! the current gate.

use crate::conv2d::Conv2d;

/// Exact DeepEncoder geometry (`deeplip_b_l`). Values are the checkpoint's, not tunables.
#[derive(Clone, Debug)]
pub struct DeepEncoderConfig {
    pub image_size: usize,   // 1024
    pub sam_patch: usize,    // 16  -> 64×64 grid
    pub sam_width: usize,    // 768
    pub sam_layers: usize,   // 12
    pub sam_heads: usize,    // 12
    pub sam_window: usize,   // 14
    pub sam_global: Vec<usize>, // [2,5,8,11]
    pub sam_mlp_ratio: f32,  // 4.0
    pub neck_channels: usize,       // 256 (SAM neck out)
    pub compress_channels: [usize; 2], // [512, 1024]
    pub clip_width: usize,   // 1024
    pub clip_layers: usize,  // 24
    pub clip_heads: usize,   // 16
    pub proj_in: usize,      // 2048 (concat SAM+CLIP)
    pub proj_out: usize,     // 1280 (LM hidden)
    pub eps: f32,            // 1e-6
}

impl Default for DeepEncoderConfig {
    fn default() -> Self {
        Self {
            image_size: 1024,
            sam_patch: 16,
            sam_width: 768,
            sam_layers: 12,
            sam_heads: 12,
            sam_window: 14,
            sam_global: vec![2, 5, 8, 11],
            sam_mlp_ratio: 4.0,
            neck_channels: 256,
            compress_channels: [512, 1024],
            clip_width: 1024,
            clip_layers: 24,
            clip_heads: 16,
            proj_in: 2048,
            proj_out: 1280,
            eps: 1e-6,
        }
    }
}

impl DeepEncoderConfig {
    /// SAM grid side after patch-embed: `image_size / sam_patch` = 64.
    pub fn grid(&self) -> usize {
        self.image_size / self.sam_patch
    }
    /// Compressed grid side after the two stride-2 convs: 64 -> 32 -> 16.
    pub fn compressed_grid(&self) -> usize {
        self.grid() / 4
    }
    /// Final vision-token count: `compressed_grid²` = 256.
    pub fn num_tokens(&self) -> usize {
        let g = self.compressed_grid();
        g * g
    }
}

// ── Preprocessing ───────────────────────────────────────────────────────────────────────────────

/// SAM/CLIP normalization. DeepSeek-OCR uses ImageNet mean/std (the SAM + CLIP convention), applied
/// after a resize to `image_size²`. Returns CHW f32 `[3, S, S]`.
pub const IMAGENET_MEAN: [f32; 3] = [0.485, 0.456, 0.406];
pub const IMAGENET_STD: [f32; 3] = [0.229, 0.224, 0.225];

/// Resize an RGB8 image to `S×S` (bicubic, PIL convention — reused from [`crate::vision`]) and
/// normalize. Returns **HWC** f32 `[S, S, 3]` — the layout [`Conv2d::forward`] consumes (it asserts
/// `[h, w, ic]`). torch's SAM `PatchEmbed` runs an NCHW conv; the numeric result is identical, only
/// the packing differs, and our Conv2d owns the packing.
pub fn preprocess(rgb8: &[u8], w: usize, h: usize, cfg: &DeepEncoderConfig) -> Vec<f32> {
    let s = cfg.image_size;
    let resized = crate::vision::resize_rgb8_bicubic(rgb8, w, h, s, s);
    let mut out = vec![0f32; s * s * 3];
    for y in 0..s {
        for x in 0..s {
            for c in 0..3 {
                let v = resized[(y * s + x) * 3 + c] as f32 / 255.0;
                out[(y * s + x) * 3 + c] = (v - IMAGENET_MEAN[c]) / IMAGENET_STD[c];
            }
        }
    }
    out
}

// ── Small reference ops (clarity over speed — this is the oracle) ─────────────────────────────────

/// Exact GELU (erf form) — SAM's MLP (`act_layer=nn.GELU`, approximate='none').
pub(crate) fn gelu(x: f32) -> f32 {
    0.5 * x * (1.0 + erf(x * std::f32::consts::FRAC_1_SQRT_2))
}
/// CLIP's `quick_gelu` = `x·sigmoid(1.702·x)` (confirmed: `deepencoder.py::quick_gelu`, used by the
/// CLIP `NoTPFeedForward`). NOT the same as SAM's exact GELU — a real numeric difference.
pub(crate) fn quick_gelu(x: f32) -> f32 {
    x / (1.0 + (-1.702 * x).exp())
}
fn erf(x: f32) -> f32 {
    // Abramowitz-Stegun 7.1.26 — reference-grade; the GPU path uses the shared clamped tanh-GELU
    // only where the checkpoint was trained with it. SAM/CLIP were not, so match erf here.
    let t = 1.0 / (1.0 + 0.3275911 * x.abs());
    let y = 1.0
        - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t
            + 0.254829592)
            * t
            * (-x * x).exp();
    if x < 0.0 { -y } else { y }
}

/// LayerNorm over the last dim `c` of a `[rows, c]` row-major buffer.
pub(crate) fn layernorm(x: &[f32], rows: usize, c: usize, w: &[f32], b: &[f32], eps: f32) -> Vec<f32> {
    let mut out = vec![0f32; rows * c];
    for r in 0..rows {
        let row = &x[r * c..r * c + c];
        let mean = row.iter().sum::<f32>() / c as f32;
        let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / c as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for j in 0..c {
            out[r * c + j] = (row[j] - mean) * inv * w[j] + b[j];
        }
    }
    out
}

/// Naive `[m,k]·[k,n]+bias -> [m,n]` (weight row-major `[n,k]`, torch Linear layout).
pub(crate) fn linear(x: &[f32], m: usize, k: usize, n: usize, w: &[f32], b: Option<&[f32]>) -> Vec<f32> {
    let mut out = vec![0f32; m * n];
    for i in 0..m {
        for j in 0..n {
            let mut acc = b.map_or(0.0, |bb| bb[j]);
            for p in 0..k {
                acc += x[i * k + p] * w[j * k + p];
            }
            out[i * n + j] = acc;
        }
    }
    out
}

fn softmax_inplace(row: &mut [f32]) {
    let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut s = 0.0;
    for v in row.iter_mut() {
        *v = (*v - m).exp();
        s += *v;
    }
    for v in row.iter_mut() {
        *v /= s;
    }
}

// ── SAM decomposed relative-position attention (the crux) ─────────────────────────────────────────

/// Weights for one SAM block (torch layouts; loaded in step 3).
pub struct SamBlockWeights {
    pub norm1_w: Vec<f32>,
    pub norm1_b: Vec<f32>,
    pub qkv_w: Vec<f32>, // [3*width, width]
    pub qkv_b: Vec<f32>, // [3*width]
    pub proj_w: Vec<f32>,
    pub proj_b: Vec<f32>,
    pub norm2_w: Vec<f32>,
    pub norm2_b: Vec<f32>,
    pub mlp_fc1_w: Vec<f32>, // [hidden, width]
    pub mlp_fc1_b: Vec<f32>,
    pub mlp_fc2_w: Vec<f32>, // [width, hidden]
    pub mlp_fc2_b: Vec<f32>,
    /// Decomposed relative-position tables `[2*max-1, head_dim]` for the height and width axes.
    pub rel_pos_h: Vec<f32>,
    pub rel_pos_w: Vec<f32>,
}

/// SAM's `get_rel_pos`: for query size `q` and key size `k`, gather (and interpolate when sizes
/// differ) the relative-position embedding of dim `hd` into `[q, k, hd]`. Here q==k (square grid or
/// window), so it's a straight gather at `q_i - k_j + (k-1)`.
pub(crate) fn get_rel_pos(q: usize, k: usize, rel: &[f32], hd: usize) -> Vec<f32> {
    // Table is [2*max_rel-1, hd]; with q==k the needed span is 2*q-1 and no interpolation is used.
    debug_assert_eq!(q, k, "reference assumes square attention (window or full grid)");
    let mut out = vec![0f32; q * k * hd];
    for i in 0..q {
        for j in 0..k {
            let idx = (i as isize - j as isize + (k as isize - 1)) as usize;
            out[(i * k + j) * hd..(i * k + j) * hd + hd]
                .copy_from_slice(&rel[idx * hd..idx * hd + hd]);
        }
    }
    out
}

/// SAM's `add_decomposed_rel_pos`: adds `rel_h + rel_w` into the `[heads, q_hw, k_hw]` attention
/// logits, where the grid is `gh×gw`. `q = einsum(r_q, Rh)` along height, similarly width.
fn add_decomposed_rel_pos(
    attn: &mut [f32], // [heads, gh*gw, gh*gw]
    q: &[f32],        // [heads, gh*gw, hd]
    heads: usize,
    gh: usize,
    gw: usize,
    hd: usize,
    rel_pos_h: &[f32],
    rel_pos_w: &[f32],
) {
    let rh = get_rel_pos(gh, gh, rel_pos_h, hd); // [gh, gh, hd]
    let rw = get_rel_pos(gw, gw, rel_pos_w, hd); // [gw, gw, hd]
    let n = gh * gw;
    for h in 0..heads {
        for qy in 0..gh {
            for qx in 0..gw {
                let qi = qy * gw + qx;
                let qvec = &q[(h * n + qi) * hd..(h * n + qi) * hd + hd];
                // Precompute dh[ky] = <qvec, Rh[qy,ky]> and dw[kx] = <qvec, Rw[qx,kx]> ONCE
                // (dw is independent of ky; recomputing it per ky was a 64× waste). Math identical.
                let dh: Vec<f32> = (0..gh)
                    .map(|ky| {
                        let rhv = &rh[(qy * gh + ky) * hd..(qy * gh + ky) * hd + hd];
                        qvec.iter().zip(rhv).map(|(a, b)| a * b).sum()
                    })
                    .collect();
                let dw: Vec<f32> = (0..gw)
                    .map(|kx| {
                        let rwv = &rw[(qx * gw + kx) * hd..(qx * gw + kx) * hd + hd];
                        qvec.iter().zip(rwv).map(|(a, b)| a * b).sum()
                    })
                    .collect();
                for ky in 0..gh {
                    for kx in 0..gw {
                        attn[(h * n + qi) * n + ky * gw + kx] += dh[ky] + dw[kx];
                    }
                }
            }
        }
    }
}

/// Multi-head attention over a `gh×gw` grid with decomposed relative position. Input/out `[n, width]`
/// (n = gh*gw). This is the shared core for both the windowed and global SAM attention (window =
/// call it on each 14×14 tile; global = call it on the whole 64×64 grid).
pub(crate) fn sam_attention(
    x: &[f32],
    gh: usize,
    gw: usize,
    width: usize,
    heads: usize,
    w: &SamBlockWeights,
) -> Vec<f32> {
    let n = gh * gw;
    let hd = width / heads;
    let scale = 1.0 / (hd as f32).sqrt();
    // qkv: [n, 3*width]
    let qkv = linear(x, n, width, 3 * width, &w.qkv_w, Some(&w.qkv_b));
    // split into per-head q,k,v: [heads, n, hd]
    let mut q = vec![0f32; heads * n * hd];
    let mut k = vec![0f32; heads * n * hd];
    let mut v = vec![0f32; heads * n * hd];
    for i in 0..n {
        for h in 0..heads {
            for d in 0..hd {
                let base = i * 3 * width;
                q[(h * n + i) * hd + d] = qkv[base + h * hd + d];
                k[(h * n + i) * hd + d] = qkv[base + width + h * hd + d];
                v[(h * n + i) * hd + d] = qkv[base + 2 * width + h * hd + d];
            }
        }
    }
    // scores [heads, n, n] = scale * q·kᵀ
    let mut attn = vec![0f32; heads * n * n];
    for h in 0..heads {
        for i in 0..n {
            for j in 0..n {
                let mut s = 0.0;
                for d in 0..hd {
                    s += q[(h * n + i) * hd + d] * k[(h * n + j) * hd + d];
                }
                attn[(h * n + i) * n + j] = s * scale;
            }
        }
    }
    // + decomposed relative position (note: torch adds rel pos BEFORE softmax, unscaled q — SAM
    // multiplies q by scale first; we fold scale into the qk product above and add rel pos on the
    // scaled q to match `attn = (q*scale) @ k.T; attn += rel`). SAM uses q (unscaled) for rel pos,
    // so pass q here, not scaled.
    add_decomposed_rel_pos(&mut attn, &q, heads, gh, gw, hd, &w.rel_pos_h, &w.rel_pos_w);
    // softmax over keys, then ·v
    let mut out = vec![0f32; heads * n * hd];
    for h in 0..heads {
        for i in 0..n {
            let row = &mut attn[(h * n + i) * n..(h * n + i) * n + n];
            softmax_inplace(row);
            for d in 0..hd {
                let mut acc = 0.0;
                for j in 0..n {
                    acc += row[j] * v[(h * n + j) * hd + d];
                }
                out[(h * n + i) * hd + d] = acc;
            }
        }
    }
    // merge heads -> [n, width]
    let mut merged = vec![0f32; n * width];
    for i in 0..n {
        for h in 0..heads {
            for d in 0..hd {
                merged[i * width + h * hd + d] = out[(h * n + i) * hd + d];
            }
        }
    }
    // output projection
    linear(&merged, n, width, width, &w.proj_w, Some(&w.proj_b))
}

/// One SAM block over the full `grid×grid` feature map `[grid*grid, width]`.
/// `windowed` selects window-14 attention (perception layers) vs full-grid global attention.
pub fn sam_block(
    x: &[f32],
    grid: usize,
    cfg: &DeepEncoderConfig,
    windowed: bool,
    w: &SamBlockWeights,
) -> Vec<f32> {
    let width = cfg.sam_width;
    let n = grid * grid;
    // pre-norm
    let normed = layernorm(x, n, width, &w.norm1_w, &w.norm1_b, cfg.eps);
    let attn_out = if windowed {
        window_attention(&normed, grid, cfg, w)
    } else {
        sam_attention(&normed, grid, grid, width, cfg.sam_heads, w)
    };
    // residual
    let mut y = vec![0f32; n * width];
    for i in 0..n * width {
        y[i] = x[i] + attn_out[i];
    }
    // MLP
    let normed2 = layernorm(&y, n, width, &w.norm2_w, &w.norm2_b, cfg.eps);
    let hidden = (width as f32 * cfg.sam_mlp_ratio) as usize;
    let mut fc1 = linear(&normed2, n, width, hidden, &w.mlp_fc1_w, Some(&w.mlp_fc1_b));
    for v in fc1.iter_mut() {
        *v = gelu(*v);
    }
    let fc2 = linear(&fc1, n, hidden, width, &w.mlp_fc2_w, Some(&w.mlp_fc2_b));
    for i in 0..n * width {
        y[i] += fc2[i];
    }
    y
}

/// SAM window-14 attention: pad the `grid×grid` map up to a multiple of `window`, partition into
/// `window×window` tiles, attend within each (with decomposed rel pos on the window), unpartition,
/// and crop back. Matches SAM's `window_partition`/`window_unpartition`.
fn window_attention(x: &[f32], grid: usize, cfg: &DeepEncoderConfig, w: &SamBlockWeights) -> Vec<f32> {
    let width = cfg.sam_width;
    let win = cfg.sam_window;
    let pad = (win - grid % win) % win;
    let gp = grid + pad; // padded side
    // pad (zeros) into [gp, gp, width]
    let mut xp = vec![0f32; gp * gp * width];
    for y in 0..grid {
        for x0 in 0..grid {
            let src = (y * grid + x0) * width;
            let dst = (y * gp + x0) * width;
            xp[dst..dst + width].copy_from_slice(&x[src..src + width]);
        }
    }
    let nw = gp / win; // windows per side
    let mut out = vec![0f32; gp * gp * width];
    for wy in 0..nw {
        for wx in 0..nw {
            // gather window [win*win, width]
            let mut tile = vec![0f32; win * win * width];
            for iy in 0..win {
                for ix in 0..win {
                    let src = ((wy * win + iy) * gp + (wx * win + ix)) * width;
                    let dst = (iy * win + ix) * width;
                    tile[dst..dst + width].copy_from_slice(&xp[src..src + width]);
                }
            }
            let att = sam_attention(&tile, win, win, width, cfg.sam_heads, w);
            for iy in 0..win {
                for ix in 0..win {
                    let dst = ((wy * win + iy) * gp + (wx * win + ix)) * width;
                    let src = (iy * win + ix) * width;
                    out[dst..dst + width].copy_from_slice(&att[src..src + width]);
                }
            }
        }
    }
    // crop back to grid×grid
    let mut cropped = vec![0f32; grid * grid * width];
    for y in 0..grid {
        for x0 in 0..grid {
            let src = (y * gp + x0) * width;
            let dst = (y * grid + x0) * width;
            cropped[dst..dst + width].copy_from_slice(&out[src..src + width]);
        }
    }
    cropped
}

// ── PatchEmbed (implemented; reuses the verified Conv2d) ──────────────────────────────────────────

/// SAM patch embedding: conv 16×16 stride 16, 3→768. Input HWC `[S,S,3]` (from [`preprocess`]),
/// output `[grid*grid, 768]` token-major over the 64×64 grid — ready for the first SAM block.
///
/// `Conv2d::forward` returns `[oh*ow, oc]` row-major (its `gemm_packed` writes `[m=oh·ow, n=oc]`),
/// which is already the token-major layout the transformer wants, so no transpose is needed.
pub fn patch_embed(hwc: &[f32], cfg: &DeepEncoderConfig, conv_w: &[f32], conv_b: &[f32]) -> Vec<f32> {
    let s = cfg.image_size;
    let conv = Conv2d::from_torch(
        conv_w,
        Some(conv_b),
        cfg.sam_width,
        3,
        cfg.sam_patch,
        cfg.sam_patch,
        cfg.sam_patch,
        0,
    );
    let (feat, oh, ow) = conv.forward(hwc, s, s);
    debug_assert_eq!(oh, cfg.grid());
    debug_assert_eq!(ow, cfg.grid());
    debug_assert_eq!(feat.len(), cfg.grid() * cfg.grid() * cfg.sam_width);
    feat
}

// ── SAM neck (768→256) ────────────────────────────────────────────────────────────────────────
//
// Standard SAM neck, applied to the 64×64×768 SAM output: conv1×1 768→256 (+LayerNorm2d),
// conv3×3 pad1 256→256 (+LayerNorm2d). LayerNorm2d normalizes over the channel dim per spatial
// position — identical to [`layernorm`] over the last dim of an HWC `[h*w, c]` buffer.
//
// WIRING TO VERIFY IN STEP 3: whether DeepSeek-OCR keeps SAM's two-LN neck verbatim or a variant,
// and the exact tensor names. The channel path (768→256) is fixed by `downsample_channels[0]=512`
// downstream requiring a 256-channel input.

/// Weights for the SAM neck (torch conv layouts).
pub struct NeckWeights {
    pub conv1_w: Vec<f32>, // [256, 768, 1, 1]
    pub ln1_w: Vec<f32>,
    pub ln1_b: Vec<f32>,
    pub conv2_w: Vec<f32>, // [256, 256, 3, 3]
    pub ln2_w: Vec<f32>,
    pub ln2_b: Vec<f32>,
}

/// SAM neck: `[grid*grid, 768]` token-major → `[grid*grid, 256]`. Input/output HWC over the grid.
pub fn sam_neck(x: &[f32], cfg: &DeepEncoderConfig, w: &NeckWeights) -> Vec<f32> {
    let g = cfg.grid();
    let cin = cfg.sam_width; // 768
    let cout = cfg.neck_channels; // 256
    let conv1 = Conv2d::from_torch(&w.conv1_w, None, cout, cin, 1, 1, 1, 0);
    let (mut y, _, _) = conv1.forward(x, g, g); // [g*g, 256]
    y = layernorm(&y, g * g, cout, &w.ln1_w, &w.ln1_b, cfg.eps);
    let conv2 = Conv2d::from_torch(&w.conv2_w, None, cout, cout, 3, 3, 1, 1);
    let (mut z, _, _) = conv2.forward(&y, g, g); // [g*g, 256]
    z = layernorm(&z, g * g, cout, &w.ln2_w, &w.ln2_b, cfg.eps);
    z
}

// ── 16× convolutional compressor: net_2, net_3 (256→512→1024, 64→32→16) ────────────────────────
//
// CONFIRMED against deepencoder.py `ImageEncoderViT`:
//   self.net_2 = nn.Conv2d(256, 512,  3, stride=2, padding=1, bias=False)
//   self.net_3 = nn.Conv2d(512, 1024, 3, stride=2, padding=1, bias=False)
//   forward: x2 = net_2(neck_out); x3 = net_3(x2); return x3   ← NO activation, NO bias.
// This is the optical-compression step: 64×64×256 → 16×16×1024, i.e. 4096 → 256 tokens.

/// Weights for the compressor's two (bias-free) convs.
pub struct CompressorWeights {
    pub net2_w: Vec<f32>, // [512, 256, 3, 3]
    pub net3_w: Vec<f32>, // [1024, 512, 3, 3]
}

/// 16× compressor: `[64*64, 256]` → `[16*16, 1024]` (token-major). No bias, no activation — exactly
/// `net_3(net_2(x))`.
pub fn compress(x: &[f32], cfg: &DeepEncoderConfig, w: &CompressorWeights) -> Vec<f32> {
    let g = cfg.grid(); // 64
    let c0 = cfg.neck_channels; // 256
    let [c1, c2] = cfg.compress_channels; // 512, 1024
    let net2 = Conv2d::from_torch(&w.net2_w, None, c1, c0, 3, 3, 2, 1);
    let (y, oh, _) = net2.forward(x, g, g); // [32*32, 512]
    debug_assert_eq!(oh, g / 2);
    let net3 = Conv2d::from_torch(&w.net3_w, None, c2, c1, 3, 3, 2, 1);
    let (z, oh2, _) = net3.forward(&y, g / 2, g / 2); // [16*16, 1024]
    debug_assert_eq!(oh2, cfg.compressed_grid());
    z
}

// ── CLIP-L tower (24 blocks, global attention on the 256 compressed tokens) ────────────────────

/// Standard pre-LN global multi-head attention (no relative position, no rope) over `n` tokens of
/// width `width`. Used by the CLIP tower; reused shape for any dense ViT block.
fn mha_global(x: &[f32], n: usize, width: usize, heads: usize, qkv_w: &[f32], qkv_b: &[f32], proj_w: &[f32], proj_b: &[f32]) -> Vec<f32> {
    let hd = width / heads;
    let scale = 1.0 / (hd as f32).sqrt();
    let qkv = linear(x, n, width, 3 * width, qkv_w, Some(qkv_b));
    let mut out = vec![0f32; n * width];
    for h in 0..heads {
        // per-head q,k,v views via strided reads from qkv[i, {0,width,2width} + h*hd + d]
        let mut attn = vec![0f32; n * n];
        for i in 0..n {
            for j in 0..n {
                let mut s = 0.0;
                for d in 0..hd {
                    let qi = qkv[i * 3 * width + h * hd + d];
                    let kj = qkv[j * 3 * width + width + h * hd + d];
                    s += qi * kj;
                }
                attn[i * n + j] = s * scale;
            }
        }
        for i in 0..n {
            softmax_inplace(&mut attn[i * n..i * n + n]);
            for d in 0..hd {
                let mut acc = 0.0;
                for j in 0..n {
                    acc += attn[i * n + j] * qkv[j * 3 * width + 2 * width + h * hd + d];
                }
                out[i * width + h * hd + d] = acc;
            }
        }
    }
    linear(&out, n, width, width, proj_w, Some(proj_b))
}

/// Weights for one CLIP block.
pub struct ClipBlockWeights {
    pub norm1_w: Vec<f32>,
    pub norm1_b: Vec<f32>,
    pub qkv_w: Vec<f32>,
    pub qkv_b: Vec<f32>,
    pub proj_w: Vec<f32>,
    pub proj_b: Vec<f32>,
    pub norm2_w: Vec<f32>,
    pub norm2_b: Vec<f32>,
    pub mlp_fc1_w: Vec<f32>,
    pub mlp_fc1_b: Vec<f32>,
    pub mlp_fc2_w: Vec<f32>,
    pub mlp_fc2_b: Vec<f32>,
}

/// One CLIP pre-LN transformer block over `[n, width]`. Global attention, MLP(×4) with GELU.
pub fn clip_block(x: &[f32], n: usize, cfg: &DeepEncoderConfig, w: &ClipBlockWeights) -> Vec<f32> {
    let width = cfg.clip_width;
    let normed = layernorm(x, n, width, &w.norm1_w, &w.norm1_b, cfg.eps);
    let attn = mha_global(&normed, n, width, cfg.clip_heads, &w.qkv_w, &w.qkv_b, &w.proj_w, &w.proj_b);
    let mut y = vec![0f32; n * width];
    for i in 0..n * width {
        y[i] = x[i] + attn[i];
    }
    let normed2 = layernorm(&y, n, width, &w.norm2_w, &w.norm2_b, cfg.eps);
    // CLIP MLP hidden = ffn_hidden_size (derive from the loaded fc1 so it can't drift), quick_gelu.
    let hidden = w.mlp_fc1_b.len();
    let mut fc1 = linear(&normed2, n, width, hidden, &w.mlp_fc1_w, Some(&w.mlp_fc1_b));
    for v in fc1.iter_mut() {
        *v = quick_gelu(*v); // CLIP uses quick_gelu, NOT exact gelu
    }
    let fc2 = linear(&fc1, n, hidden, width, &w.mlp_fc2_w, Some(&w.mlp_fc2_b));
    for i in 0..n * width {
        y[i] += fc2[i];
    }
    y
}

/// CLIP tower embeddings + weights. CONFIRMED against deepencoder.py `VitModel`/`CLIPVisionEmbeddings`:
/// the tower takes the SAM compressor output as its `patch_embeds` (bypassing CLIP's own patch
/// conv), prepends a learned class token, adds position embeddings, applies `pre_layrnorm`, runs 24
/// blocks, and the caller drops the class token (`output[:, 1:]`).
pub struct ClipTowerWeights {
    pub class_embedding: Vec<f32>,      // [1024]
    pub position_embedding: Vec<f32>,   // [num_pos, 1024] (num_pos = num_tokens + 1)
    pub pre_ln_w: Vec<f32>,
    pub pre_ln_b: Vec<f32>,
    pub blocks: Vec<ClipBlockWeights>,  // 24
}

/// Run the CLIP-L tower on the SAM compressor tokens `patch_embeds` `[n, 1024]`. Returns the
/// post-transformer tokens WITHOUT the class token: `[n, 1024]`.
pub fn clip_tower(patch_embeds: &[f32], n: usize, cfg: &DeepEncoderConfig, w: &ClipTowerWeights) -> Vec<f32> {
    let width = cfg.clip_width;
    let seq = n + 1; // + class token
    // embeddings: [class ; patch_embeds] + position_embedding
    let mut x = vec![0f32; seq * width];
    x[0..width].copy_from_slice(&w.class_embedding);
    x[width..seq * width].copy_from_slice(&patch_embeds[0..n * width]);
    debug_assert_eq!(w.position_embedding.len(), seq * width, "pos-embed length must be n+1 tokens");
    for i in 0..seq * width {
        x[i] += w.position_embedding[i];
    }
    // pre_layrnorm, then the transformer stack
    let mut h = layernorm(&x, seq, width, &w.pre_ln_w, &w.pre_ln_b, cfg.eps);
    for blk in &w.blocks {
        h = clip_block(&h, seq, cfg, blk);
    }
    // drop the class token -> [n, width]
    h[width..seq * width].to_vec()
}

// ── Projector (concat CLIP+SAM → LM hidden) ───────────────────────────────────────────────────
//
// CONFIRMED against modeling_deepseekocr.py:
//   local_features = torch.cat((clip[:, 1:], sam.flatten(2).permute(0,2,1)), dim=-1)  # CLIP, then SAM
//   local_features = self.projector(local_features)                                   # linear 2048->1280

/// Projector: `linear(concat(clip_i, sam_i)) -> [n, proj_out]`. `clip` and `sam` are both
/// `[n, 1024]` (CLIP tower output and the SAM compressor output). CLIP is the FIRST half.
pub fn project(clip: &[f32], sam: &[f32], n: usize, cfg: &DeepEncoderConfig, w: &[f32], bias: &[f32]) -> Vec<f32> {
    let half = cfg.clip_width; // 1024
    debug_assert_eq!(cfg.proj_in, 2 * half);
    let mut concat = vec![0f32; n * cfg.proj_in];
    for i in 0..n {
        concat[i * cfg.proj_in..i * cfg.proj_in + half].copy_from_slice(&clip[i * half..i * half + half]);
        concat[i * cfg.proj_in + half..(i + 1) * cfg.proj_in].copy_from_slice(&sam[i * half..i * half + half]);
    }
    linear(&concat, n, cfg.proj_in, cfg.proj_out, w, Some(bias))
}

// ── Full assembly: weights, loader, forward ───────────────────────────────────────────────────

/// All DeepEncoder weights, loaded from a `deepseek-ai/DeepSeek-OCR` checkpoint directory.
pub struct DeepEncoderWeights {
    pub patch_embed_w: Vec<f32>, // [768, 3, 16, 16]
    pub patch_embed_b: Vec<f32>,
    pub sam_pos_embed: Vec<f32>, // [1, 64, 64, 768] -> flattened [grid*grid, 768]
    pub sam_blocks: Vec<SamBlockWeights>, // 12
    pub neck: NeckWeights,
    pub compressor: CompressorWeights,
    pub clip: ClipTowerWeights,
    pub proj_w: Vec<f32>, // [1280, 2048]
    pub proj_b: Vec<f32>,
}

/// The exact tensor-name manifest this encoder loads, generated from `cfg`. Verified against the
/// checkpoint's `model.safetensors.index.json`. `load` and the structural-completeness gate share
/// this so the two can never drift.
pub fn tensor_manifest(cfg: &DeepEncoderConfig) -> Vec<String> {
    let mut m = Vec::new();
    let sam = "model.sam_model";
    m.push(format!("{sam}.patch_embed.proj.weight"));
    m.push(format!("{sam}.patch_embed.proj.bias"));
    m.push(format!("{sam}.pos_embed"));
    for i in 0..cfg.sam_layers {
        let b = format!("{sam}.blocks.{i}");
        for t in [
            "norm1.weight", "norm1.bias", "attn.qkv.weight", "attn.qkv.bias",
            "attn.proj.weight", "attn.proj.bias", "attn.rel_pos_h", "attn.rel_pos_w",
            "norm2.weight", "norm2.bias", "mlp.lin1.weight", "mlp.lin1.bias",
            "mlp.lin2.weight", "mlp.lin2.bias",
        ] {
            m.push(format!("{b}.{t}"));
        }
    }
    // neck.0 conv, neck.1 LN, neck.2 conv, neck.3 LN (convs bias-free)
    m.push(format!("{sam}.neck.0.weight"));
    m.push(format!("{sam}.neck.1.weight"));
    m.push(format!("{sam}.neck.1.bias"));
    m.push(format!("{sam}.neck.2.weight"));
    m.push(format!("{sam}.neck.3.weight"));
    m.push(format!("{sam}.neck.3.bias"));
    m.push(format!("{sam}.net_2.weight"));
    m.push(format!("{sam}.net_3.weight"));
    let clip = "model.vision_model";
    m.push(format!("{clip}.embeddings.class_embedding"));
    m.push(format!("{clip}.embeddings.position_embedding.weight"));
    m.push(format!("{clip}.pre_layrnorm.weight"));
    m.push(format!("{clip}.pre_layrnorm.bias"));
    for i in 0..cfg.clip_layers {
        let b = format!("{clip}.transformer.layers.{i}");
        for t in [
            "layer_norm1.weight", "layer_norm1.bias", "self_attn.qkv_proj.weight",
            "self_attn.qkv_proj.bias", "self_attn.out_proj.weight", "self_attn.out_proj.bias",
            "layer_norm2.weight", "layer_norm2.bias", "mlp.fc1.weight", "mlp.fc1.bias",
            "mlp.fc2.weight", "mlp.fc2.bias",
        ] {
            m.push(format!("{b}.{t}"));
        }
    }
    m.push("model.projector.layers.weight".into());
    m.push("model.projector.layers.bias".into());
    m
}

impl DeepEncoderWeights {
    /// Load from a checkpoint directory (single or sharded safetensors). Uses [`tensor_manifest`]
    /// names. NOTE: the numeric correctness of the result is gated by the torch reference (step 3c);
    /// this only guarantees the tensors exist and shapes are consistent.
    pub fn load(dir: &std::path::Path, cfg: &DeepEncoderConfig) -> anyhow::Result<Self> {
        let st = crate::weights::LazySt::open(dir)?;
        let g = |n: &str| st.tensor_f32(n);
        let sam = "model.sam_model";
        let sam_block = |i: usize| -> anyhow::Result<SamBlockWeights> {
            let b = format!("{sam}.blocks.{i}");
            Ok(SamBlockWeights {
                norm1_w: g(&format!("{b}.norm1.weight"))?, norm1_b: g(&format!("{b}.norm1.bias"))?,
                qkv_w: g(&format!("{b}.attn.qkv.weight"))?, qkv_b: g(&format!("{b}.attn.qkv.bias"))?,
                proj_w: g(&format!("{b}.attn.proj.weight"))?, proj_b: g(&format!("{b}.attn.proj.bias"))?,
                norm2_w: g(&format!("{b}.norm2.weight"))?, norm2_b: g(&format!("{b}.norm2.bias"))?,
                mlp_fc1_w: g(&format!("{b}.mlp.lin1.weight"))?, mlp_fc1_b: g(&format!("{b}.mlp.lin1.bias"))?,
                mlp_fc2_w: g(&format!("{b}.mlp.lin2.weight"))?, mlp_fc2_b: g(&format!("{b}.mlp.lin2.bias"))?,
                rel_pos_h: g(&format!("{b}.attn.rel_pos_h"))?, rel_pos_w: g(&format!("{b}.attn.rel_pos_w"))?,
            })
        };
        let clip = "model.vision_model";
        let clip_block = |i: usize| -> anyhow::Result<ClipBlockWeights> {
            let b = format!("{clip}.transformer.layers.{i}");
            Ok(ClipBlockWeights {
                norm1_w: g(&format!("{b}.layer_norm1.weight"))?, norm1_b: g(&format!("{b}.layer_norm1.bias"))?,
                qkv_w: g(&format!("{b}.self_attn.qkv_proj.weight"))?, qkv_b: g(&format!("{b}.self_attn.qkv_proj.bias"))?,
                proj_w: g(&format!("{b}.self_attn.out_proj.weight"))?, proj_b: g(&format!("{b}.self_attn.out_proj.bias"))?,
                norm2_w: g(&format!("{b}.layer_norm2.weight"))?, norm2_b: g(&format!("{b}.layer_norm2.bias"))?,
                mlp_fc1_w: g(&format!("{b}.mlp.fc1.weight"))?, mlp_fc1_b: g(&format!("{b}.mlp.fc1.bias"))?,
                mlp_fc2_w: g(&format!("{b}.mlp.fc2.weight"))?, mlp_fc2_b: g(&format!("{b}.mlp.fc2.bias"))?,
            })
        };
        Ok(Self {
            patch_embed_w: g(&format!("{sam}.patch_embed.proj.weight"))?,
            patch_embed_b: g(&format!("{sam}.patch_embed.proj.bias"))?,
            sam_pos_embed: g(&format!("{sam}.pos_embed"))?,
            sam_blocks: (0..cfg.sam_layers).map(sam_block).collect::<anyhow::Result<_>>()?,
            neck: NeckWeights {
                conv1_w: g(&format!("{sam}.neck.0.weight"))?,
                ln1_w: g(&format!("{sam}.neck.1.weight"))?, ln1_b: g(&format!("{sam}.neck.1.bias"))?,
                conv2_w: g(&format!("{sam}.neck.2.weight"))?,
                ln2_w: g(&format!("{sam}.neck.3.weight"))?, ln2_b: g(&format!("{sam}.neck.3.bias"))?,
            },
            compressor: CompressorWeights {
                net2_w: g(&format!("{sam}.net_2.weight"))?,
                net3_w: g(&format!("{sam}.net_3.weight"))?,
            },
            clip: ClipTowerWeights {
                class_embedding: g(&format!("{clip}.embeddings.class_embedding"))?,
                position_embedding: g(&format!("{clip}.embeddings.position_embedding.weight"))?,
                pre_ln_w: g(&format!("{clip}.pre_layrnorm.weight"))?,
                pre_ln_b: g(&format!("{clip}.pre_layrnorm.bias"))?,
                blocks: (0..cfg.clip_layers).map(clip_block).collect::<anyhow::Result<_>>()?,
            },
            proj_w: g("model.projector.layers.weight")?,
            proj_b: g("model.projector.layers.bias")?,
        })
    }
}

/// Full DeepEncoder forward: RGB8 image → `[256, 1280]` vision tokens, matching
/// modeling_deepseekocr.py's data flow. This is the reference the GPU path (step 4) is gated on.
pub fn forward(rgb8: &[u8], w: usize, h: usize, cfg: &DeepEncoderConfig, wt: &DeepEncoderWeights) -> Vec<f32> {
    let g = cfg.grid();
    let width = cfg.sam_width;
    // preprocess + patch embed -> [g*g, 768]
    let hwc = preprocess(rgb8, w, h, cfg);
    let mut x = patch_embed(&hwc, cfg, &wt.patch_embed_w, &wt.patch_embed_b);
    // + absolute position embedding (SAM pos_embed [1,g,g,768] -> [g*g,768])
    for i in 0..g * g * width {
        x[i] += wt.sam_pos_embed[i];
    }
    // 12 SAM blocks: global at cfg.sam_global, windowed elsewhere
    for (i, blk) in wt.sam_blocks.iter().enumerate() {
        let global = cfg.sam_global.contains(&i);
        x = sam_block(&x, g, cfg, !global, blk);
    }
    // neck 768->256, compressor 256->1024 @ 16×16 -> the 256 SAM tokens
    let necked = sam_neck(&x, cfg, &wt.neck);
    let sam_tokens = compress(&necked, cfg, &wt.compressor); // [256, 1024]
    let n = cfg.num_tokens();
    // CLIP tower on the SAM tokens -> [256, 1024]
    let clip_tokens = clip_tower(&sam_tokens, n, cfg, &wt.clip);
    // projector: concat(clip, sam) 2048 -> 1280
    project(&clip_tokens, &sam_tokens, n, cfg, &wt.proj_w, &wt.proj_b)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn geometry_is_exact() {
        let c = DeepEncoderConfig::default();
        assert_eq!(c.grid(), 64);
        assert_eq!(c.compressed_grid(), 16);
        assert_eq!(c.num_tokens(), 256);
    }

    #[test]
    fn preprocess_shape_and_norm() {
        let c = DeepEncoderConfig::default();
        // 2×2 white image -> resized to 1024², normalized. Channel 0 of white: (1-mean)/std.
        let rgb = vec![255u8; 2 * 2 * 3];
        let chw = preprocess(&rgb, 2, 2, &c);
        assert_eq!(chw.len(), 3 * 1024 * 1024);
        let expect0 = (1.0 - IMAGENET_MEAN[0]) / IMAGENET_STD[0];
        assert!((chw[0] - expect0).abs() < 1e-3, "got {}", chw[0]);
    }

    #[test]
    fn patch_embed_grid_shape() {
        let c = DeepEncoderConfig::default();
        let chw = vec![0.1f32; 3 * c.image_size * c.image_size];
        // identity-ish conv weights [oc=768, ic=3, 16,16], zero bias
        let conv_w = vec![0.001f32; 768 * 3 * 16 * 16];
        let conv_b = vec![0f32; 768];
        let tokens = patch_embed(&chw, &c, &conv_w, &conv_b);
        assert_eq!(tokens.len(), 64 * 64 * 768);
    }

    #[test]
    fn sam_block_preserves_shape() {
        // Tiny grid to keep the O(n²) reference cheap: 4×4 grid, width 24, 4 heads.
        let mut c = DeepEncoderConfig::default();
        c.sam_width = 24;
        c.sam_heads = 4;
        c.sam_window = 4;
        let grid = 4;
        let n = grid * grid;
        let width = c.sam_width;
        let hd = width / c.sam_heads;
        let x = vec![0.02f32; n * width];
        let w = SamBlockWeights {
            norm1_w: vec![1.0; width], norm1_b: vec![0.0; width],
            qkv_w: vec![0.01; 3 * width * width], qkv_b: vec![0.0; 3 * width],
            proj_w: vec![0.01; width * width], proj_b: vec![0.0; width],
            norm2_w: vec![1.0; width], norm2_b: vec![0.0; width],
            mlp_fc1_w: vec![0.01; width * 4 * width], mlp_fc1_b: vec![0.0; 4 * width],
            mlp_fc2_w: vec![0.01; width * 4 * width], mlp_fc2_b: vec![0.0; width],
            rel_pos_h: vec![0.0; (2 * grid - 1) * hd],
            rel_pos_w: vec![0.0; (2 * grid - 1) * hd],
        };
        let global = sam_block(&x, grid, &c, false, &w);
        assert_eq!(global.len(), n * width);
        let windowed = sam_block(&x, grid, &c, true, &w);
        assert_eq!(windowed.len(), n * width);
        assert!(global.iter().all(|v| v.is_finite()));
        assert!(windowed.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn compressor_halves_twice() {
        let c = DeepEncoderConfig::default();
        let g = c.grid(); // 64
        let x = vec![0.01f32; g * g * c.neck_channels];
        let w = CompressorWeights {
            net2_w: vec![0.001; 512 * 256 * 3 * 3],
            net3_w: vec![0.001; 1024 * 512 * 3 * 3],
        };
        let out = compress(&x, &c, &w);
        // 16×16×1024 = 256 tokens × 1024
        assert_eq!(out.len(), c.compressed_grid() * c.compressed_grid() * c.compress_channels[1]);
        assert_eq!(out.len(), c.num_tokens() * 1024);
    }

    #[test]
    fn neck_768_to_256() {
        let c = DeepEncoderConfig::default();
        let g = c.grid();
        let x = vec![0.01f32; g * g * c.sam_width];
        let w = NeckWeights {
            conv1_w: vec![0.001; 256 * 768], ln1_w: vec![1.0; 256], ln1_b: vec![0.0; 256],
            conv2_w: vec![0.001; 256 * 256 * 9], ln2_w: vec![1.0; 256], ln2_b: vec![0.0; 256],
        };
        let out = sam_neck(&x, &c, &w);
        assert_eq!(out.len(), g * g * c.neck_channels);
        assert!(out.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn clip_block_and_projector_shapes() {
        let c = DeepEncoderConfig::default();
        let n = 16; // small token count for the O(n²) reference
        let width = c.clip_width;
        let x = vec![0.005f32; n * width];
        let w = ClipBlockWeights {
            norm1_w: vec![1.0; width], norm1_b: vec![0.0; width],
            qkv_w: vec![0.001; 3 * width * width], qkv_b: vec![0.0; 3 * width],
            proj_w: vec![0.001; width * width], proj_b: vec![0.0; width],
            norm2_w: vec![1.0; width], norm2_b: vec![0.0; width],
            mlp_fc1_w: vec![0.001; 4 * width * width], mlp_fc1_b: vec![0.0; 4 * width],
            mlp_fc2_w: vec![0.001; 4 * width * width], mlp_fc2_b: vec![0.0; width],
        };
        let out = clip_block(&x, n, &c, &w);
        assert_eq!(out.len(), n * width);
        assert!(out.iter().all(|v| v.is_finite()));
        // projector: concat(clip 1024, sam 1024)=2048 -> 1280
        let clipf = vec![0.1f32; n * width];
        let samf = vec![0.2f32; n * width];
        let pw = vec![0.0005f32; c.proj_out * c.proj_in];
        let pb = vec![0.0f32; c.proj_out];
        let proj = project(&clipf, &samf, n, &c, &pw, &pb);
        assert_eq!(proj.len(), n * c.proj_out);
        // concat order is CLIP-first: first half of each concat row is the clip value (0.1).
        // (checked indirectly: projector shape + the doc-confirmed order.)
    }

    #[test]
    fn tensor_manifest_is_complete_and_exact() {
        let c = DeepEncoderConfig::default();
        let m = tensor_manifest(&c);
        // exact count: patch(2)+pos(1)+12*14 sam-block + neck(6) + net_2/3(2)
        //            + clip embed(4) + 24*12 clip-block + proj(2)
        let expected = 2 + 1 + 12 * 14 + 6 + 2 + 4 + 24 * 12 + 2;
        assert_eq!(m.len(), expected, "manifest count drift");
        // spot-check the exact real names confirmed from the checkpoint index
        for name in [
            "model.sam_model.patch_embed.proj.weight",
            "model.sam_model.pos_embed",
            "model.sam_model.blocks.11.attn.rel_pos_w",
            "model.sam_model.neck.0.weight",
            "model.sam_model.net_2.weight",
            "model.sam_model.net_3.weight",
            "model.vision_model.embeddings.class_embedding",
            "model.vision_model.pre_layrnorm.weight",
            "model.vision_model.transformer.layers.23.self_attn.qkv_proj.weight",
            "model.projector.layers.weight",
        ] {
            assert!(m.contains(&name.to_string()), "manifest missing {name}");
        }
        // no duplicates
        let mut sorted = m.clone();
        sorted.sort();
        sorted.dedup();
        assert_eq!(sorted.len(), m.len(), "duplicate tensor names in manifest");
    }

    #[test]
    fn clip_tower_drops_class_token() {
        let mut c = DeepEncoderConfig::default();
        c.clip_width = 32;
        c.clip_heads = 4;
        let n = 8;
        let width = c.clip_width;
        let seq = n + 1;
        let mk_block = || ClipBlockWeights {
            norm1_w: vec![1.0; width], norm1_b: vec![0.0; width],
            qkv_w: vec![0.001; 3 * width * width], qkv_b: vec![0.0; 3 * width],
            proj_w: vec![0.001; width * width], proj_b: vec![0.0; width],
            norm2_w: vec![1.0; width], norm2_b: vec![0.0; width],
            mlp_fc1_w: vec![0.001; 4 * width * width], mlp_fc1_b: vec![0.0; 4 * width],
            mlp_fc2_w: vec![0.001; 4 * width * width], mlp_fc2_b: vec![0.0; width],
        };
        let tw = ClipTowerWeights {
            class_embedding: vec![0.01; width],
            position_embedding: vec![0.0; seq * width],
            pre_ln_w: vec![1.0; width], pre_ln_b: vec![0.0; width],
            blocks: (0..2).map(|_| mk_block()).collect(),
        };
        let patch = vec![0.02f32; n * width];
        let out = clip_tower(&patch, n, &c, &tw);
        // output has the class token dropped: n tokens, not n+1
        assert_eq!(out.len(), n * width);
        assert!(out.iter().all(|v| v.is_finite()));
    }
}