//! Running TensorFlow Lite models without TensorFlow: a reader for the `.tflite` FlatBuffer
//! (fields as numbered in TensorFlow Lite's published `schema.fbs`, Apache-2.0) and an
//! interpreter for the float operators small vision models use: `CONV_2D`, `DEPTHWISE_CONV_2D`,
//! `MAX_POOL_2D`, `ADD`, `PAD`, `CONCATENATION`, `RESHAPE`, `RELU`, `RELU6`, `PRELU`, `LOGISTIC`
//! and `DEQUANTIZE` of float16 weights (folded at load). Weights are re-laid out once at load for
//! the [`nn`](crate::nn) kernels. Anything else (other operators, quantised tensors, dilation) is a
//! load error, never a panic: models are pinned by SHA-256, so what loads once always loads.

use crate::nn::{Conv, Depthwise, Linear};
use crate::pt::half_to_f32;

type Result<T> = std::result::Result<T, String>;

// ---------------------------------------------------------------- FlatBuffers

/// A FlatBuffer: bounds-checked little-endian reads.
#[derive(Clone, Copy)]
struct Fb<'a>(&'a [u8]);

/// A table: its position in the buffer.
#[derive(Clone, Copy)]
struct Table(usize);

impl<'a> Fb<'a> {
    fn bytes<const N: usize>(&self, at: usize) -> Result<[u8; N]> {
        self.0.get(at..at.checked_add(N).ok_or("bad offset")?).and_then(|s| s.try_into().ok()).ok_or_else(|| "truncated model".to_string())
    }
    fn u8(&self, at: usize) -> Result<u8> {
        Ok(self.bytes::<1>(at)?[0])
    }
    fn u16(&self, at: usize) -> Result<u16> {
        Ok(u16::from_le_bytes(self.bytes(at)?))
    }
    fn u32(&self, at: usize) -> Result<u32> {
        Ok(u32::from_le_bytes(self.bytes(at)?))
    }
    fn i32(&self, at: usize) -> Result<i32> {
        Ok(i32::from_le_bytes(self.bytes(at)?))
    }
    fn u64(&self, at: usize) -> Result<u64> {
        Ok(u64::from_le_bytes(self.bytes(at)?))
    }
    /// Follow the offset stored at `at`.
    fn deref(&self, at: usize) -> Result<usize> {
        at.checked_add(self.u32(at)? as usize).ok_or_else(|| "bad offset".to_string())
    }
    fn root(&self) -> Result<Table> {
        Ok(Table(self.deref(0)?))
    }
    /// Where field `id` of `t` is stored (absent fields have their default).
    fn field(&self, t: Table, id: usize) -> Result<Option<usize>> {
        let vt = (t.0 as i64 - self.i32(t.0)? as i64).try_into().map_err(|_| "bad vtable")?;
        let vt_len = self.u16(vt)? as usize;
        let slot = 4 + 2 * id;
        if slot + 2 > vt_len {
            return Ok(None);
        }
        let off = self.u16(vt + slot)? as usize;
        Ok((off != 0).then_some(t.0 + off))
    }
    fn int(&self, t: Table, id: usize, default: i32) -> Result<i32> {
        self.field(t, id)?.map_or(Ok(default), |p| self.i32(p))
    }
    fn byte(&self, t: Table, id: usize, default: u8) -> Result<u8> {
        self.field(t, id)?.map_or(Ok(default), |p| self.u8(p))
    }
    fn table(&self, t: Table, id: usize) -> Result<Option<Table>> {
        self.field(t, id)?.map(|p| self.deref(p).map(Table)).transpose()
    }
    /// A vector field: (start of its elements, length).
    fn vector(&self, t: Table, id: usize) -> Result<(usize, usize)> {
        let Some(p) = self.field(t, id)? else { return Ok((0, 0)) };
        let v = self.deref(p)?;
        let n = self.u32(v)? as usize;
        if v.saturating_add(4).saturating_add(n) > self.0.len() {
            return Err("truncated model".into());
        }
        Ok((v + 4, n))
    }
    fn tables(&self, t: Table, id: usize) -> Result<Vec<Table>> {
        let (at, n) = self.vector(t, id)?;
        (0..n).map(|i| self.deref(at + 4 * i).map(Table)).collect()
    }
    fn ints(&self, t: Table, id: usize) -> Result<Vec<i32>> {
        let (at, n) = self.vector(t, id)?;
        (0..n).map(|i| self.i32(at + 4 * i)).collect()
    }
    fn string(&self, t: Table, id: usize) -> Result<String> {
        let (at, n) = self.vector(t, id)?;
        Ok(String::from_utf8_lossy(self.0.get(at..at + n).ok_or("truncated model")?).into_owned())
    }
}

// ---------------------------------------------------------------- the model

/// Operator codes (`BuiltinOperator`).
mod op {
    pub const ADD: i32 = 0;
    pub const CONCATENATION: i32 = 2;
    pub const CONV_2D: i32 = 3;
    pub const DEPTHWISE_CONV_2D: i32 = 4;
    pub const DEQUANTIZE: i32 = 6;
    pub const LOGISTIC: i32 = 14;
    pub const MAX_POOL_2D: i32 = 17;
    pub const RELU: i32 = 19;
    pub const RELU6: i32 = 21;
    pub const RESHAPE: i32 = 22;
    pub const PAD: i32 = 34;
    pub const PRELU: i32 = 54;
}

/// The most values a tensor may hold.
const MAX_TENSOR: usize = 1 << 28;

/// Tensor types (`TensorType`).
const FLOAT32: u8 = 0;
const FLOAT16: u8 = 1;
const INT32: u8 = 2;

/// A fused activation (`ActivationFunctionType`).
#[derive(Clone, Copy, Debug, PartialEq)]
enum Act {
    None,
    Relu,
    ReluN1To1,
    Relu6,
    Tanh,
}

impl Act {
    fn from(code: u8) -> Result<Act> {
        Ok(match code {
            0 => Act::None,
            1 => Act::Relu,
            2 => Act::ReluN1To1,
            3 => Act::Relu6,
            4 => Act::Tanh,
            c => return Err(format!("unsupported fused activation {c}")),
        })
    }
    fn apply(self, x: &mut [f32]) {
        let f: fn(f32) -> f32 = match self {
            Act::None => return,
            Act::Relu => |v| v.max(0.0),
            Act::ReluN1To1 => |v| v.clamp(-1.0, 1.0),
            Act::Relu6 => |v| v.clamp(0.0, 6.0),
            Act::Tanh => f32::tanh,
        };
        for v in x {
            *v = f(*v);
        }
    }
}

/// One operator, its weights laid out for the kernels.
#[derive(Clone, Debug)]
enum Op {
    /// A 1×1, stride-1 convolution: a dense layer over the pixels.
    Pointwise(Linear, Act),
    Conv(Conv, [usize; 2], Act),
    Depthwise(Depthwise, [usize; 2], Act),
    /// Kernel, stride, padding before (top, left).
    MaxPool {
        k: [usize; 2],
        stride: [usize; 2],
        pad: [usize; 2],
        act: Act,
    },
    Add(Act),
    /// Constant padding (zeros), per dimension (before, after).
    Pad(Vec<[usize; 2]>),
    Concat {
        axis: usize,
        act: Act,
    },
    Reshape,
    Act(Act),
    Prelu(Vec<f32>),
    Logistic,
}

#[derive(Clone, Debug)]
struct Step {
    op: Op,
    /// Runtime inputs (tensor ids).
    inputs: Vec<usize>,
    output: usize,
}

/// A loaded model, ready to run (batch of one, static shapes).
#[derive(Clone, Debug)]
pub struct Model {
    shapes: Vec<Vec<usize>>,
    names: Vec<String>,
    /// Constant float tensors (weights folded into the steps are not kept).
    consts: Vec<Option<Vec<f32>>>,
    steps: Vec<Step>,
    inputs: Vec<usize>,
    outputs: Vec<usize>,
    /// After which step each tensor is last needed (to free memory as the model runs).
    last_use: Vec<usize>,
}

fn index(i: i32, n: usize) -> Result<usize> {
    usize::try_from(i).ok().filter(|&i| i < n).ok_or_else(|| format!("bad tensor index {i}"))
}

/// "Same" padding: output size and padding before.
fn same(input: usize, k: usize, stride: usize) -> (usize, usize) {
    let out = input.div_ceil(stride);
    let total = ((out.saturating_sub(1)) * stride + k).saturating_sub(input);
    (out, total / 2)
}

/// Output size and padding before, for `padding` (0 same, 1 valid).
fn window(padding: u8, input: usize, k: usize, stride: usize) -> Result<(usize, usize)> {
    match padding {
        0 => Ok(same(input, k, stride)),
        1 if input >= k => Ok(((input - k) / stride + 1, 0)),
        _ => Err("bad window".into()),
    }
}

impl Model {
    /// Read a `.tflite` file.
    pub fn read(bytes: &[u8]) -> Result<Model> {
        let fb = Fb(bytes);
        let root = fb.root()?;
        let codes: Vec<i32> = fb.tables(root, 1)?.into_iter().map(|c| Ok(fb.int(c, 3, 0)?.max(fb.byte(c, 0, 0)? as i8 as i32))).collect::<Result<_>>()?;
        let buffers = fb.tables(root, 4)?;
        let graphs = fb.tables(root, 2)?;
        let g = *graphs.first().ok_or("the model has no graph")?;
        let tensors = fb.tables(g, 0)?;
        let n = tensors.len();
        let mut shapes = Vec::with_capacity(n);
        let mut names = Vec::with_capacity(n);
        // Constant data: float tensors as f32, int tensors (pads, shapes) as i32.
        let mut floats: Vec<Option<Vec<f32>>> = vec![None; n];
        let mut ints: Vec<Option<Vec<i32>>> = vec![None; n];
        for (i, t) in tensors.iter().enumerate() {
            let shape: Vec<usize> = fb
                .ints(*t, 0)?
                .into_iter()
                .map(|d| usize::try_from(d).map_err(|_| "dynamic shapes are not supported"))
                .collect::<std::result::Result<_, _>>()?;
            let ty = fb.byte(*t, 1, FLOAT32)?;
            let buf = fb.int(*t, 2, 0)? as u32 as usize;
            names.push(fb.string(*t, 3)?);
            // Every size computed from shapes later stays in range.
            let count = shape.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).filter(|&c| c <= MAX_TENSOR).ok_or("a tensor is too large")?;
            if fb.field(*t, 4)?.is_some() && fb.table(*t, 4)?.is_some_and(|q| fb.vector(q, 2).is_ok_and(|v| v.1 > 0)) {
                return Err(format!("{}: quantised tensors are not supported", names[i]));
            }
            if buf > 0 {
                let b = *buffers.get(buf).ok_or("bad buffer index")?;
                let (mut at, mut len) = fb.vector(b, 0)?;
                if len == 0 {
                    // Large models keep data after the FlatBuffer: offset and size from its start.
                    let (o, s) = (fb.field(b, 1)?.map(|p| fb.u64(p)).transpose()?.unwrap_or(0), fb.field(b, 2)?.map(|p| fb.u64(p)).transpose()?.unwrap_or(0));
                    (at, len) = (o as usize, s as usize);
                }
                if len > 0 {
                    let data = bytes.get(at..at.checked_add(len).ok_or("bad buffer")?).ok_or("truncated model")?;
                    match ty {
                        FLOAT32 if data.len() == 4 * count => {
                            floats[i] = Some(data.as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
                        }
                        FLOAT16 if data.len() == 2 * count => {
                            floats[i] = Some(data.as_chunks::<2>().0.iter().map(|c| half_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect())
                        }
                        INT32 if data.len() == 4 * count => {
                            ints[i] = Some(data.as_chunks::<4>().0.iter().map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
                        }
                        _ => return Err(format!("{}: unsupported tensor data", names[i])),
                    }
                }
            }
            shapes.push(shape);
        }
        let inputs = fb.ints(g, 1)?.into_iter().map(|i| index(i, n)).collect::<Result<Vec<_>>>()?;
        let outputs = fb.ints(g, 2)?.into_iter().map(|i| index(i, n)).collect::<Result<Vec<_>>>()?;
        let mut steps = vec![];
        for o in fb.tables(g, 3)? {
            let code = *codes.get(fb.int(o, 0, 0)? as u32 as usize).ok_or("bad operator code")?;
            // Optional inputs are -1.
            let ins: Vec<Option<usize>> = fb.ints(o, 1)?.into_iter().map(|i| if i < 0 { Ok(None) } else { index(i, n).map(Some) }).collect::<Result<_>>()?;
            let outs = fb.ints(o, 2)?;
            let output = index(*outs.first().ok_or("an operator has no output")?, n)?;
            // Float16 weights: fold their DEQUANTIZE.
            if code == op::DEQUANTIZE
                && let Some(f) = ins.first().copied().flatten().and_then(|i| floats[i].clone())
            {
                floats[output] = Some(f);
                continue;
            }
            let opts = fb.table(o, 4)?;
            let opt_int = |id: usize, default: i32| opts.map_or(Ok(default), |t| fb.int(t, id, default));
            let opt_byte = |id: usize, default: u8| opts.map_or(Ok(default), |t| fb.byte(t, id, default));
            let input = |k: usize| ins.get(k).copied().flatten().ok_or_else(|| format!("{}: missing input {k}", names[output]));
            let weights = |k: usize| -> Result<Vec<f32>> {
                let id = input(k)?;
                floats[id].clone().ok_or_else(|| format!("{}: input {k} is not constant", names[output]))
            };
            let bias = |count: usize| -> Result<Vec<f32>> {
                match ins.get(2).copied().flatten() {
                    Some(id) => floats[id].clone().filter(|b| b.len() == count).ok_or_else(|| format!("{}: bad bias", names[output])),
                    None => Ok(vec![0.0; count]),
                }
            };
            let in_shape = |k: usize| -> Result<Vec<usize>> { Ok(shapes[input(k)?].clone()) };
            let out_shape = shapes[output].clone();
            let stride = |w: usize, h: usize| -> Result<[usize; 2]> {
                let s = [opt_int(h, 1)?, opt_int(w, 1)?];
                if s.iter().any(|v| *v < 1) {
                    return Err("bad stride".into());
                }
                Ok(s.map(|v| v as usize))
            };
            let mut runtime = vec![input(0)?];
            let op = match code {
                op::DEQUANTIZE => Op::Reshape,
                op::CONV_2D | op::DEPTHWISE_CONV_2D => {
                    let depthwise = code == op::DEPTHWISE_CONV_2D;
                    let (padding, s) = (opt_byte(0, 0)?, stride(1, 2)?);
                    let act_id = if depthwise { 4 } else { 3 };
                    let dil = if depthwise { [opt_int(6, 1)?, opt_int(5, 1)?] } else { [opt_int(5, 1)?, opt_int(4, 1)?] };
                    if dil != [1, 1] {
                        return Err(format!("{}: dilated convolutions are not supported", names[output]));
                    }
                    let act = Act::from(opt_byte(act_id, 0)?)?;
                    let (x, f) = (in_shape(0)?, shapes[input(1)?].clone());
                    let (&[1, h, w, cin], &[fo, kh, kw, fi], &[1, oh, ow, cout]) = (&x[..], &f[..], &out_shape[..]) else {
                        return Err(format!("{}: unexpected convolution shapes", names[output]));
                    };
                    if kh != kw || s[0] != s[1] {
                        return Err(format!("{}: only square kernels and strides are supported", names[output]));
                    }
                    let (ey, py) = window(padding, h, kh, s[0])?;
                    let (ex, px) = window(padding, w, kw, s[1])?;
                    if (ey, ex) != (oh, ow) {
                        return Err(format!("{}: output size disagrees with the padding", names[output]));
                    }
                    let wt = weights(1)?;
                    let b = bias(cout)?;
                    if depthwise {
                        if fo != 1 || fi != cout || cin != cout {
                            return Err(format!("{}: depth multipliers are not supported", names[output]));
                        }
                        Op::Depthwise(Depthwise { w: wt, b, c: cout, ks: kh, stride: s[0] }, [py, px], act)
                    } else {
                        if fo != cout || fi != cin {
                            return Err(format!("{}: unexpected filter shape", names[output]));
                        }
                        // OHWI → [(ky·ks + kx)·inp + ci] × out.
                        let kk = kh * kw * cin;
                        let mut t = vec![0.0f32; kk * cout];
                        for o in 0..cout {
                            for p in 0..kk {
                                t[p * cout + o] = *wt.get(o * kk + p).ok_or("bad filter")?;
                            }
                        }
                        if kh == 1 && s[0] == 1 {
                            Op::Pointwise(Linear { wt: t, b, inp: cin, out: cout }, act)
                        } else {
                            Op::Conv(Conv { wt: t, b, inp: cin, out: cout, ks: kh, stride: s[0], pad: 0 }, [py, px], act)
                        }
                    }
                }
                op::MAX_POOL_2D => {
                    let (padding, s) = (opt_byte(0, 0)?, stride(1, 2)?);
                    let k = [opt_int(4, 1)?, opt_int(3, 1)?];
                    if k.iter().any(|v| *v < 1) {
                        return Err("bad pool size".into());
                    }
                    let k = k.map(|v| v as usize);
                    let x = in_shape(0)?;
                    let (&[1, h, w, _], &[1, oh, ow, _]) = (&x[..], &out_shape[..]) else { return Err("unexpected pool shapes".into()) };
                    let (ey, py) = window(padding, h, k[0], s[0])?;
                    let (ex, px) = window(padding, w, k[1], s[1])?;
                    if (ey, ex) != (oh, ow) {
                        return Err(format!("{}: output size disagrees with the padding", names[output]));
                    }
                    Op::MaxPool { k, stride: s, pad: [py, px], act: Act::from(opt_byte(5, 0)?)? }
                }
                op::ADD => {
                    runtime.push(input(1)?);
                    Op::Add(Act::from(opt_byte(0, 0)?)?)
                }
                op::PAD => {
                    let p = ints[input(1)?].clone().ok_or("PAD: paddings must be constant")?;
                    let pads = p
                        .as_chunks::<2>()
                        .0
                        .iter()
                        .map(|c| Ok([usize::try_from(c[0]).map_err(|_| "bad pad")?, usize::try_from(c[1]).map_err(|_| "bad pad")?]))
                        .collect::<Result<Vec<_>>>()?;
                    Op::Pad(pads)
                }
                op::CONCATENATION => {
                    runtime = (0..ins.len()).map(input).collect::<Result<_>>()?;
                    let rank = out_shape.len() as i32;
                    let a = opt_int(0, 0)?;
                    let axis = if a < 0 { a + rank } else { a };
                    Op::Concat { axis: index(axis, out_shape.len())?, act: Act::from(opt_byte(1, 0)?)? }
                }
                op::RESHAPE => Op::Reshape,
                op::RELU => Op::Act(Act::Relu),
                op::RELU6 => Op::Act(Act::Relu6),
                op::LOGISTIC => Op::Logistic,
                op::PRELU => {
                    let alpha = weights(1)?;
                    let c = *out_shape.last().ok_or("PRELU: bad shape")?;
                    if alpha.len() != c {
                        return Err(format!("{}: only per-channel PReLU is supported", names[output]));
                    }
                    Op::Prelu(alpha)
                }
                c => return Err(format!("operator {c} is not supported")),
            };
            steps.push(Step { op, inputs: runtime, output });
        }
        let mut last_use = vec![usize::MAX; n];
        for (k, s) in steps.iter().enumerate() {
            for &i in &s.inputs {
                last_use[i] = k;
            }
        }
        for &o in &outputs {
            last_use[o] = usize::MAX;
        }
        Ok(Model { shapes, names, consts: floats, steps, inputs, outputs, last_use })
    }

    /// The shape of input `k`.
    pub fn input_shape(&self, k: usize) -> Option<&[usize]> {
        self.inputs.get(k).and_then(|&i| self.shapes.get(i)).map(Vec::as_slice)
    }

    /// The shape and name of output `k`.
    pub fn output(&self, k: usize) -> Option<(&[usize], &str)> {
        let &i = self.outputs.get(k)?;
        Some((self.shapes.get(i)?.as_slice(), self.names.get(i)?.as_str()))
    }

    /// Run the model on its inputs (row-major, NHWC); returns its outputs.
    pub fn run(&self, inputs: &[&[f32]]) -> Result<Vec<Vec<f32>>> {
        let mut vals: Vec<Option<Vec<f32>>> = vec![None; self.shapes.len()];
        for (k, &id) in self.inputs.iter().enumerate() {
            let x = inputs.get(k).ok_or("missing model input")?;
            if x.len() != self.shapes[id].iter().product::<usize>() {
                return Err(format!("input {k}: expected {:?}", self.shapes[id]));
            }
            vals[id] = Some(x.to_vec());
        }
        for (k, s) in self.steps.iter().enumerate() {
            let y = {
                let get = |i: usize| -> Result<&[f32]> {
                    vals.get(i)
                        .and_then(Option::as_deref)
                        .or_else(|| self.consts.get(i).and_then(Option::as_deref))
                        .ok_or_else(|| format!("{}: input not computed", self.names[s.output]))
                };
                let x = get(s.inputs[0])?;
                let xs = &self.shapes[s.inputs[0]];
                let ys = &self.shapes[s.output];
                let (h, w) = (xs.get(1).copied().unwrap_or(1), xs.get(2).copied().unwrap_or(1));
                let (oh, ow) = (ys.get(1).copied().unwrap_or(1), ys.get(2).copied().unwrap_or(1));
                let y = match &s.op {
                    Op::Pointwise(l, a) => {
                        let mut y = l.forward(x, h * w);
                        a.apply(&mut y);
                        y
                    }
                    Op::Conv(c, pad, a) => {
                        let mut y = c.forward_sized(x, h, w, oh, ow, *pad);
                        a.apply(&mut y);
                        y
                    }
                    Op::Depthwise(d, pad, a) => {
                        let mut y = d.forward_sized(x, h, w, oh, ow, *pad);
                        a.apply(&mut y);
                        y
                    }
                    Op::MaxPool { k, stride, pad, act } => {
                        let mut y = max_pool(x, [h, w], *xs.last().unwrap_or(&1), [oh, ow], *k, *stride, *pad);
                        act.apply(&mut y);
                        y
                    }
                    Op::Add(a) => {
                        let b = get(s.inputs[1])?;
                        if b.len() != x.len() {
                            return Err(format!("{}: broadcasting ADD is not supported", self.names[s.output]));
                        }
                        let mut y: Vec<f32> = x.iter().zip(b).map(|(p, q)| p + q).collect();
                        a.apply(&mut y);
                        y
                    }
                    Op::Pad(p) => pad(x, xs, p)?,
                    Op::Concat { axis, act } => {
                        let parts = s.inputs.iter().map(|&i| Ok((get(i)?, &self.shapes[i]))).collect::<Result<Vec<_>>>()?;
                        let mut y = concat(&parts, *axis)?;
                        act.apply(&mut y);
                        y
                    }
                    Op::Reshape => x.to_vec(),
                    Op::Act(a) => {
                        let mut y = x.to_vec();
                        a.apply(&mut y);
                        y
                    }
                    Op::Prelu(alpha) => {
                        let c = alpha.len().max(1);
                        x.iter().enumerate().map(|(i, &v)| if v >= 0.0 { v } else { v * alpha[i % c] }).collect()
                    }
                    Op::Logistic => x.iter().map(|v| 1.0 / (1.0 + (-v).exp())).collect(),
                };
                if y.len() != ys.iter().product::<usize>() {
                    return Err(format!("{}: computed {} values for shape {ys:?}", self.names[s.output], y.len()));
                }
                y
            };
            vals[s.output] = Some(y);
            for &i in &s.inputs {
                if self.last_use[i] == k {
                    vals[i] = None;
                }
            }
        }
        self.outputs.iter().map(|&o| vals[o].take().or_else(|| self.consts[o].clone()).ok_or_else(|| "an output was not computed".to_string())).collect()
    }
}

/// Max pooling over `k` windows; window cells outside the input are ignored.
fn max_pool(x: &[f32], [h, w]: [usize; 2], c: usize, [oh, ow]: [usize; 2], k: [usize; 2], s: [usize; 2], pad: [usize; 2]) -> Vec<f32> {
    let mut y = vec![f32::NEG_INFINITY; oh * ow * c];
    for oy in 0..oh {
        for ox in 0..ow {
            let out = &mut y[(oy * ow + ox) * c..][..c];
            for ky in 0..k[0] {
                let Some(iy) = (oy * s[0] + ky).checked_sub(pad[0]).filter(|&v| v < h) else { continue };
                for kx in 0..k[1] {
                    let Some(ix) = (ox * s[1] + kx).checked_sub(pad[1]).filter(|&v| v < w) else { continue };
                    let Some(src) = x.get((iy * w + ix) * c..(iy * w + ix + 1) * c) else { continue };
                    for (o, v) in out.iter_mut().zip(src) {
                        *o = o.max(*v);
                    }
                }
            }
        }
    }
    y
}

/// Zero padding of a row-major tensor of shape `shape`.
fn pad(x: &[f32], shape: &[usize], pads: &[[usize; 2]]) -> Result<Vec<f32>> {
    if pads.len() != shape.len() {
        return Err("PAD: rank mismatch".into());
    }
    let out: Vec<usize> = shape.iter().zip(pads).map(|(d, p)| d + p[0] + p[1]).collect();
    let mut y = vec![0.0f32; out.iter().product()];
    let rank = shape.len();
    let inner = *shape.last().unwrap_or(&1);
    let rows = x.len() / inner.max(1);
    // Copy each innermost row to its padded place.
    let mut idx = vec![0usize; rank.saturating_sub(1)];
    for r in 0..rows {
        let mut at = 0usize;
        for d in 0..rank {
            let i = if d + 1 == rank { pads[d][0] } else { idx[d] + pads[d][0] };
            at = at * out[d] + i;
        }
        if let (Some(dst), Some(src)) = (y.get_mut(at..at + inner), x.get(r * inner..(r + 1) * inner)) {
            dst.copy_from_slice(src);
        }
        for d in (0..idx.len()).rev() {
            idx[d] += 1;
            if idx[d] < shape[d] {
                break;
            }
            idx[d] = 0;
        }
    }
    Ok(y)
}

/// Concatenate along `axis`.
fn concat(parts: &[(&[f32], &Vec<usize>)], axis: usize) -> Result<Vec<f32>> {
    let Some((_, s0)) = parts.first() else { return Ok(vec![]) };
    let outer: usize = s0.get(..axis).ok_or("CONCATENATION: bad axis")?.iter().product();
    let mut y = Vec::with_capacity(parts.iter().map(|p| p.0.len()).sum());
    for o in 0..outer {
        for (x, s) in parts {
            let chunk: usize = s.get(axis..).ok_or("CONCATENATION: bad axis")?.iter().product();
            y.extend_from_slice(x.get(o * chunk..(o + 1) * chunk).ok_or("CONCATENATION: bad input")?);
        }
    }
    Ok(y)
}

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

    #[test]
    fn same_padding_matches_tensorflow() {
        // TensorFlow pads the extra row after: 128 → 64 with a 3×3 stride-2 window pads 0 before.
        assert_eq!(same(128, 3, 2), (64, 0));
        assert_eq!(same(128, 5, 2), (64, 1));
        assert_eq!(same(16, 3, 1), (16, 1));
        assert_eq!(same(7, 2, 2), (4, 0));
        assert_eq!(window(1, 4, 2, 2), Ok((2, 0)));
        assert!(window(1, 1, 2, 2).is_err());
    }

    #[test]
    fn asymmetric_conv_padding() {
        // A 3×3 stride-2 "same" conv on 4×4: pads nothing before, one after.
        let x: Vec<f32> = (0..16).map(|v| v as f32).collect();
        let ones = vec![1.0f32; 9];
        let c = Conv { wt: ones.clone(), b: vec![0.0], inp: 1, out: 1, ks: 3, stride: 2, pad: 0 };
        let y = c.forward_sized(&x, 4, 4, 2, 2, [0, 0]);
        // Top-left window rows 0–2, cols 0–2; the bottom-right one is cut by the edge.
        assert_eq!(y, vec![45.0, 39.0, 66.0, 50.0]);
        let d = Depthwise { w: ones, b: vec![0.0], c: 1, ks: 3, stride: 2 };
        assert_eq!(d.forward_sized(&x, 4, 4, 2, 2, [0, 0]), y);
    }

    #[test]
    fn pool_pad_and_concat() {
        let x: Vec<f32> = (0..16).map(|v| v as f32).collect();
        assert_eq!(max_pool(&x, [4, 4], 1, [2, 2], [2, 2], [2, 2], [0, 0]), vec![5.0, 7.0, 13.0, 15.0]);
        // Same padding on 3×3 → 2×2: the last window is cut by the edge.
        let x3: Vec<f32> = (0..9).map(|v| v as f32).collect();
        assert_eq!(max_pool(&x3, [3, 3], 1, [2, 2], [2, 2], [2, 2], [0, 0]), vec![4.0, 5.0, 7.0, 8.0]);
        // Channel padding (residual connections widen the channels with zeros).
        let p = pad(&[1.0, 2.0, 3.0, 4.0], &[1, 1, 2, 2], &[[0, 0], [0, 0], [0, 0], [0, 1]]).unwrap();
        assert_eq!(p, vec![1.0, 2.0, 0.0, 3.0, 4.0, 0.0]);
        let p = pad(&[1.0, 2.0], &[1, 2], &[[0, 0], [1, 0]]).unwrap();
        assert_eq!(p, vec![0.0, 1.0, 2.0]);
        let a = [1.0, 2.0, 3.0, 4.0];
        let b = [5.0, 6.0];
        let (sa, sb) = (vec![1, 2, 2], vec![1, 1, 2]);
        assert_eq!(concat(&[(&a, &sa), (&b, &sb)], 1).unwrap(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let (sa, sb) = (vec![2, 2], vec![2, 1]);
        assert_eq!(concat(&[(&a, &sa), (&b, &sb)], 1).unwrap(), vec![1.0, 2.0, 5.0, 3.0, 4.0, 6.0]);
    }

    #[test]
    fn damaged_files_are_errors_not_panics() {
        assert!(Model::read(b"").is_err());
        assert!(Model::read(&[0xff; 64]).is_err());
        // Every prefix and a few corruptions of a hand-made buffer.
        let mut fb = vec![12, 0, 0, 0, 0, 0, 6, 0, 8, 0, 4, 0, 6, 0, 0, 0, 0, 0, 0, 0];
        for n in 0..fb.len() {
            let _ = Model::read(&fb[..n]);
        }
        for i in 0..fb.len() {
            fb[i] ^= 0x5a;
            let _ = Model::read(&fb);
        }
    }
}