//! The few neural-network kernels MobileSAM needs, on token-major feature maps (`[pixels, channels]`
//! row-major, i.e. NHWC with a batch of one): dense layers on a parallel GEMM (ndarray's safe
//! `general_mat_mul`, which picks AVX/FMA kernels at run time), convolutions as patch matrices
//! times weights, depthwise and transposed convolutions, layer norm, exact GELU and softmax.
//! Weights are re-laid out once at load so every inner loop runs over contiguous channels.

use ndarray::linalg::general_mat_mul;
use ndarray::{ArrayView2, ArrayViewMut2};
use rayon::prelude::*;

/// Rows of the left matrix per parallel task.
const ROW_BLOCK: usize = 64;

/// `out[m×n] += a[m×k] · b[k×n]` (row-major), rows split across threads.
pub fn gemm(a: &[f32], m: usize, k: usize, b: &[f32], n: usize, out: &mut [f32]) {
    if m == 0 || n == 0 || a.len() < m * k || b.len() < k * n || out.len() < m * n {
        return;
    }
    let Ok(bv) = ArrayView2::from_shape((k, n), &b[..k * n]) else { return };
    let block = |r0: usize, oc: &mut [f32]| {
        let rows = oc.len() / n;
        if let (Ok(av), Ok(mut cv)) = (ArrayView2::from_shape((rows, k), &a[r0 * k..(r0 + rows) * k]), ArrayViewMut2::from_shape((rows, n), oc)) {
            general_mat_mul(1.0, &av, &bv, 1.0, &mut cv);
        }
    };
    if m <= ROW_BLOCK {
        block(0, &mut out[..m * n]);
    } else {
        out[..m * n].par_chunks_mut(ROW_BLOCK * n).enumerate().for_each(|(i, oc)| block(i * ROW_BLOCK, oc));
    }
}

/// A dense layer (also a 1×1 convolution): `y = x·Wᵀ + b`, with `W` stored transposed.
#[derive(Clone, Debug, Default)]
pub struct Linear {
    /// `[inp × out]`.
    pub wt: Vec<f32>,
    pub b: Vec<f32>,
    pub inp: usize,
    pub out: usize,
}

impl Linear {
    /// From PyTorch's `[out × inp]` weight (a 1×1 conv's `[out × inp × 1 × 1]` too).
    pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize) -> Option<Linear> {
        if w.len() != out * inp || b.is_some_and(|b| b.len() != out) {
            return None;
        }
        let mut wt = vec![0.0; inp * out];
        for o in 0..out {
            for i in 0..inp {
                wt[i * out + o] = w[o * inp + i];
            }
        }
        Some(Linear { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out })
    }

    /// `rows` tokens of `inp` channels → `rows × out`.
    pub fn forward(&self, x: &[f32], rows: usize) -> Vec<f32> {
        let mut y: Vec<f32> = (0..rows).flat_map(|_| self.b.iter().copied()).collect();
        gemm(x, rows, self.inp, &self.wt, self.out, &mut y);
        y
    }
}

/// A `ks×ks` convolution (any stride and padding, one group) on `h×w×inp` → `oh×ow×out`, as a
/// patch matrix times the weights, built a block of output rows at a time.
#[derive(Clone, Debug, Default)]
pub struct Conv {
    /// `[(ky·ks + kx)·inp + ci] × out`.
    pub wt: Vec<f32>,
    pub b: Vec<f32>,
    pub inp: usize,
    pub out: usize,
    pub ks: usize,
    pub stride: usize,
    pub pad: usize,
}

impl Conv {
    /// From PyTorch's `[out × inp × ks × ks]`.
    pub fn new(w: &[f32], b: Option<&[f32]>, out: usize, inp: usize, ks: usize, stride: usize, pad: usize) -> Option<Conv> {
        if w.len() != out * inp * ks * ks || b.is_some_and(|b| b.len() != out) || stride == 0 {
            return None;
        }
        let kk = ks * ks * inp;
        let mut wt = vec![0.0; kk * out];
        for o in 0..out {
            for ci in 0..inp {
                for t in 0..ks * ks {
                    wt[(t * inp + ci) * out + o] = w[(o * inp + ci) * ks * ks + t];
                }
            }
        }
        Some(Conv { wt, b: b.map_or_else(|| vec![0.0; out], <[f32]>::to_vec), inp, out, ks, stride, pad })
    }

    pub fn out_size(&self, n: usize) -> usize {
        (n + 2 * self.pad).saturating_sub(self.ks) / self.stride + 1
    }

    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (oh, ow) = (self.out_size(h), self.out_size(w));
        (self.forward_sized(x, h, w, oh, ow, [self.pad; 2]), oh, ow)
    }

    /// An `oh×ow` output with `pad` (top, left) zero rows and columns before the input; what
    /// falls past the input's far edges is zero too (so asymmetric "same" padding works).
    pub fn forward_sized(&self, x: &[f32], h: usize, w: usize, oh: usize, ow: usize, pad: [usize; 2]) -> Vec<f32> {
        let kk = self.ks * self.ks * self.inp;
        let mut y: Vec<f32> = (0..oh * ow).flat_map(|_| self.b.iter().copied()).collect();
        if x.len() < h * w * self.inp || self.wt.len() < kk * self.out {
            return y;
        }
        y.par_chunks_mut(ow * self.out).enumerate().for_each(|(oy, row)| {
            // This output row's patches.
            let mut patch = vec![0.0f32; ow * kk];
            for ox in 0..ow {
                for ky in 0..self.ks {
                    let iy = (oy * self.stride + ky) as isize - pad[0] as isize;
                    if iy < 0 || iy >= h as isize {
                        continue;
                    }
                    for kx in 0..self.ks {
                        let ix = (ox * self.stride + kx) as isize - pad[1] as isize;
                        if ix < 0 || ix >= w as isize {
                            continue;
                        }
                        let src = (iy as usize * w + ix as usize) * self.inp;
                        let dst = ox * kk + (ky * self.ks + kx) * self.inp;
                        patch[dst..dst + self.inp].copy_from_slice(&x[src..src + self.inp]);
                    }
                }
            }
            if let (Ok(av), Ok(bv), Ok(mut cv)) = (
                ArrayView2::from_shape((ow, kk), &patch[..]),
                ArrayView2::from_shape((kk, self.out), &self.wt[..]),
                ArrayViewMut2::from_shape((ow, self.out), row),
            ) {
                general_mat_mul(1.0, &av, &bv, 1.0, &mut cv);
            }
        });
        y
    }
}

/// A depthwise `ks×ks` convolution (one filter per channel) on `h×w×c`.
#[derive(Clone, Debug, Default)]
pub struct Depthwise {
    /// `[tap × c]`.
    pub w: Vec<f32>,
    pub b: Vec<f32>,
    pub c: usize,
    pub ks: usize,
    pub stride: usize,
}

impl Depthwise {
    /// From PyTorch's `[c × 1 × 3 × 3]`.
    pub fn new(w: &[f32], b: Option<&[f32]>, c: usize, stride: usize) -> Option<Depthwise> {
        if w.len() != c * 9 || b.is_some_and(|b| b.len() != c) || stride == 0 {
            return None;
        }
        let mut wt = vec![0.0; 9 * c];
        for ch in 0..c {
            for t in 0..9 {
                wt[t * c + ch] = w[ch * 9 + t];
            }
        }
        Some(Depthwise { w: wt, b: b.map_or_else(|| vec![0.0; c], <[f32]>::to_vec), c, ks: 3, stride })
    }

    /// 3×3 with padding 1.
    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        let (oh, ow) = ((h + 2 - 3) / self.stride + 1, (w + 2 - 3) / self.stride + 1);
        (self.forward_sized(x, h, w, oh, ow, [1, 1]), oh, ow)
    }

    /// An `oh×ow` output with `pad` (top, left) before the input, as [`Conv::forward_sized`].
    pub fn forward_sized(&self, x: &[f32], h: usize, w: usize, oh: usize, ow: usize, pad: [usize; 2]) -> Vec<f32> {
        let (c, ks) = (self.c, self.ks);
        let mut y = vec![0.0f32; oh * ow * c];
        if x.len() < h * w * c || self.w.len() < ks * ks * c || self.b.len() < c {
            return y;
        }
        y.par_chunks_mut(ow * c).enumerate().for_each(|(oy, row)| {
            for ox in 0..ow {
                let out = &mut row[ox * c..(ox + 1) * c];
                out.copy_from_slice(&self.b);
                for ky in 0..ks {
                    let iy = (oy * self.stride + ky) as isize - pad[0] as isize;
                    if iy < 0 || iy >= h as isize {
                        continue;
                    }
                    for kx in 0..ks {
                        let ix = (ox * self.stride + kx) as isize - pad[1] as isize;
                        if ix < 0 || ix >= w as isize {
                            continue;
                        }
                        let src = &x[(iy as usize * w + ix as usize) * c..][..c];
                        let tap = &self.w[(ky * ks + kx) * c..][..c];
                        for ((o, s), k) in out.iter_mut().zip(src).zip(tap) {
                            *o += s * k;
                        }
                    }
                }
            }
        });
        y
    }
}

/// A 2×2, stride-2 transposed convolution (each input pixel paints a 2×2 block):
/// `h×w×inp` → `2h×2w×out`.
#[derive(Clone, Debug, Default)]
pub struct Upconv {
    /// One dense layer per output tap (dy, dx).
    pub taps: Vec<Linear>,
    pub out: usize,
}

impl Upconv {
    /// From PyTorch's ConvTranspose2d `[inp × out × 2 × 2]`.
    pub fn new(w: &[f32], b: &[f32], inp: usize, out: usize) -> Option<Upconv> {
        if w.len() != inp * out * 4 || b.len() != out {
            return None;
        }
        let taps = (0..4)
            .map(|t| {
                let mut wt = vec![0.0; inp * out];
                for i in 0..inp {
                    for o in 0..out {
                        wt[i * out + o] = w[(i * out + o) * 4 + t];
                    }
                }
                Linear { wt, b: b.to_vec(), inp, out }
            })
            .collect();
        Some(Upconv { taps, out })
    }

    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
        let (ow, c) = (2 * w, self.out);
        let mut y = vec![0.0f32; 4 * h * w * c];
        for (t, lin) in self.taps.iter().enumerate() {
            let (dy, dx) = (t / 2, t % 2);
            let part = lin.forward(x, h * w);
            for p in 0..h * w {
                let (py, px) = (p / w, p % w);
                let dst = ((2 * py + dy) * ow + 2 * px + dx) * c;
                if let (Some(d), Some(s)) = (y.get_mut(dst..dst + c), part.get(p * c..(p + 1) * c)) {
                    d.copy_from_slice(s);
                }
            }
        }
        y
    }
}

/// Layer norm over the channels of each token.
#[derive(Clone, Debug, Default)]
pub struct Norm {
    pub w: Vec<f32>,
    pub b: Vec<f32>,
    pub eps: f32,
}

impl Norm {
    pub fn forward(&self, x: &mut [f32]) {
        let c = self.w.len();
        if c == 0 {
            return;
        }
        x.par_chunks_mut(c).for_each(|t| {
            let mean = t.iter().sum::<f32>() / c as f32;
            let var = t.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / c as f32;
            let inv = 1.0 / (var + self.eps).sqrt();
            for ((v, w), b) in t.iter_mut().zip(&self.w).zip(&self.b) {
                *v = (*v - mean) * inv * w + b;
            }
        });
    }
}

/// The error function (Abramowitz & Stegun 7.1.26 refined: |error| < 1.5e-7).
fn erf(x: f32) -> f32 {
    let s = x.signum();
    let x = x.abs() as f64;
    let t = 1.0 / (1.0 + 0.327_591_1 * x);
    let y = 1.0 - (((((1.061_405_429 * t - 1.453_152_027) * t) + 1.421_413_741) * t - 0.284_496_736) * t + 0.254_829_592) * t * (-x * x).exp();
    s * y as f32
}

/// GELU (the exact, erf form, as PyTorch's `nn.GELU`), in place.
pub fn gelu(x: &mut [f32]) {
    x.par_chunks_mut(4096).for_each(|c| {
        for v in c {
            *v = 0.5 * *v * (1.0 + erf(*v * std::f32::consts::FRAC_1_SQRT_2));
        }
    });
}

pub fn relu(x: &mut [f32]) {
    for v in x {
        *v = v.max(0.0);
    }
}

/// `a += b`.
pub fn add(a: &mut [f32], b: &[f32]) {
    a.par_chunks_mut(4096).zip(b.par_chunks(4096)).for_each(|(a, b)| {
        for (x, y) in a.iter_mut().zip(b) {
            *x += y;
        }
    });
}

/// Softmax of each row of `n` values, in place.
pub fn softmax_rows(x: &mut [f32], n: usize) {
    if n == 0 {
        return;
    }
    for row in x.chunks_mut(n) {
        let m = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let mut s = 0.0;
        for v in row.iter_mut() {
            *v = (*v - m).exp();
            s += *v;
        }
        let inv = 1.0 / s.max(1e-30);
        for v in row.iter_mut() {
            *v *= inv;
        }
    }
}

/// Multi-head scaled dot-product attention of `nq` queries over `nk` keys/values, each `heads ×
/// hd` wide (token-major, heads side by side), plus an optional per-head bias `[heads × nq × nk]`.
pub fn attention(q: &[f32], k: &[f32], v: &[f32], nq: usize, nk: usize, heads: usize, hd: usize, bias: Option<&[f32]>) -> Vec<f32> {
    let dim = heads * hd;
    let mut out = vec![0.0f32; nq * dim];
    if q.len() < nq * dim || k.len() < nk * dim || v.len() < nk * dim {
        return out;
    }
    let scale = 1.0 / (hd as f32).sqrt();
    for h in 0..heads {
        // This head's q, kᵀ and v as contiguous matrices.
        let qh: Vec<f32> = (0..nq).flat_map(|i| q[i * dim + h * hd..][..hd].iter().map(|x| x * scale)).collect();
        let mut kt = vec![0.0f32; hd * nk];
        for j in 0..nk {
            for d in 0..hd {
                kt[d * nk + j] = k[j * dim + h * hd + d];
            }
        }
        let vh: Vec<f32> = (0..nk).flat_map(|j| v[j * dim + h * hd..][..hd].iter().copied()).collect();
        let mut s = match bias.and_then(|b| b.get(h * nq * nk..(h + 1) * nq * nk)) {
            Some(b) => b.to_vec(),
            None => vec![0.0; nq * nk],
        };
        gemm(&qh, nq, hd, &kt, nk, &mut s);
        softmax_rows(&mut s, nk);
        let mut oh = vec![0.0f32; nq * hd];
        gemm(&s, nq, nk, &vh, hd, &mut oh);
        for i in 0..nq {
            out[i * dim + h * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]);
        }
    }
    out
}

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

    #[test]
    fn gemm_matches_naive() {
        let (m, k, n) = (130, 17, 9);
        let a: Vec<f32> = (0..m * k).map(|i| ((i * 7) % 13) as f32 - 6.0).collect();
        let b: Vec<f32> = (0..k * n).map(|i| ((i * 5) % 11) as f32 - 5.0).collect();
        let mut c = vec![1.0f32; m * n];
        gemm(&a, m, k, &b, n, &mut c);
        for i in 0..m {
            for j in 0..n {
                let want: f32 = 1.0 + (0..k).map(|p| a[i * k + p] * b[p * n + j]).sum::<f32>();
                assert_eq!(c[i * n + j], want);
            }
        }
    }

    #[test]
    fn convolutions_match_direct_sums() {
        // 3×3 stride-2 conv on a 5×4×2 input vs. the definition.
        let (h, w, inp, out) = (5, 4, 2, 3);
        let x: Vec<f32> = (0..h * w * inp).map(|i| (i % 7) as f32 * 0.5).collect();
        let wt: Vec<f32> = (0..out * inp * 9).map(|i| ((i % 5) as f32 - 2.0) * 0.25).collect();
        let b = [0.1, -0.2, 0.3];
        let conv = Conv::new(&wt, Some(&b), out, inp, 3, 2, 1).unwrap();
        let (y, oh, ow) = conv.forward(&x, h, w);
        assert_eq!((oh, ow), (3, 2));
        for oy in 0..oh {
            for ox in 0..ow {
                for o in 0..out {
                    let mut s = b[o];
                    for ci in 0..inp {
                        for ky in 0..3 {
                            for kx in 0..3 {
                                let (iy, ix) = ((oy * 2 + ky) as isize - 1, (ox * 2 + kx) as isize - 1);
                                if iy >= 0 && ix >= 0 && iy < h as isize && ix < w as isize {
                                    s += x[(iy as usize * w + ix as usize) * inp + ci] * wt[((o * inp + ci) * 3 + ky) * 3 + kx];
                                }
                            }
                        }
                    }
                    assert!((y[(oy * ow + ox) * out + o] - s).abs() < 1e-5);
                }
            }
        }
        // Depthwise == a grouped conv: compare with Conv on each channel alone.
        let dw_w: Vec<f32> = (0..inp * 9).map(|i| (i as f32 - 8.0) * 0.1).collect();
        let dw = Depthwise::new(&dw_w, None, inp, 1).unwrap();
        let (yd, _, _) = dw.forward(&x, h, w);
        for ci in 0..inp {
            let xc: Vec<f32> = (0..h * w).map(|p| x[p * inp + ci]).collect();
            let single = Conv::new(&dw_w[ci * 9..ci * 9 + 9], None, 1, 1, 3, 1, 1).unwrap();
            let (yc, _, _) = single.forward(&xc, h, w);
            for p in 0..h * w {
                assert!((yd[p * inp + ci] - yc[p]).abs() < 1e-5);
            }
        }
    }

    #[test]
    fn gelu_and_softmax() {
        let mut x = [-3.0f32, -1.0, 0.0, 1.0, 3.0];
        gelu(&mut x);
        // PyTorch's exact GELU values.
        let want = [-0.004_049_7, -0.158_655_3, 0.0, 0.841_344_7, 2.995_950_3];
        for (a, b) in x.iter().zip(want) {
            assert!((a - b).abs() < 1e-5, "{a} {b}");
        }
        let mut s = [1.0f32, 2.0, 3.0];
        softmax_rows(&mut s, 3);
        assert!((s.iter().sum::<f32>() - 1.0).abs() < 1e-6 && s[2] > s[1]);
    }
}