brain2qwerty 0.0.1

//! Lightweight f32 tensors for inference, I/O, and the pure-Rust reference ops.

use std::fmt;

fn erf_f32(x: f32) -> f32 {
    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x = x.abs();
    let t = 1.0 / (1.0 + 0.3275911 * x);
    let poly = t
        * (0.254829592
            + t * (-0.284496736 + t * (1.421413741 + t * (-1.453152027 + t * 1.061405429))));
    sign * (1.0 - poly * (-x * x).exp())
}

#[derive(Clone)]
pub struct Tensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

impl fmt::Debug for Tensor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Tensor({:?}, len={})", self.shape, self.data.len())
    }
}

impl Tensor {
    pub fn zeros(shape: &[usize]) -> Self {
        let n: usize = shape.iter().product();
        Self {
            data: vec![0.0; n],
            shape: shape.to_vec(),
        }
    }

    pub fn from_vec(data: Vec<f32>, shape: Vec<usize>) -> Self {
        let n: usize = shape.iter().product();
        assert_eq!(data.len(), n);
        Self { data, shape }
    }

    pub fn numel(&self) -> usize {
        self.data.len()
    }

    pub fn strides(&self) -> Vec<usize> {
        let mut s = vec![1usize; self.shape.len()];
        for i in (0..self.shape.len().saturating_sub(1)).rev() {
            s[i] = s[i + 1] * self.shape[i + 1];
        }
        s
    }

    fn index_at(&self, idx: &[usize]) -> usize {
        let st = self.strides();
        idx.iter().enumerate().map(|(i, &v)| v * st[i]).sum()
    }

    pub fn reshape(&self, new_shape: &[usize]) -> Self {
        Self::from_vec(self.data.clone(), new_shape.to_vec())
    }

    pub fn transpose(&self, axes: &[usize]) -> Self {
        let nd = self.shape.len();
        assert_eq!(axes.len(), nd);
        let new_shape: Vec<usize> = axes.iter().map(|&a| self.shape[a]).collect();
        let out = Self::zeros(&new_shape);
        let mut idx = vec![0usize; nd];
        let mut out_data = out.data.clone();
        loop {
            let src: Vec<usize> = (0..nd).map(|i| idx[axes[i]]).collect();
            let dst: Vec<usize> = idx.clone();
            out_data[out.index_at(&dst)] = self.data[self.index_at(&src)];
            if !Self::bump_indices(&mut idx, &new_shape) {
                break;
            }
        }
        Self {
            data: out_data,
            shape: new_shape,
        }
    }

    fn bump_indices(idx: &mut [usize], shape: &[usize]) -> bool {
        for i in (0..idx.len()).rev() {
            idx[i] += 1;
            if idx[i] < shape[i] {
                return true;
            }
            idx[i] = 0;
        }
        false
    }

    pub fn narrow(&self, dim: usize, start: usize, length: usize) -> Self {
        assert!(dim < self.shape.len());
        let mut new_shape = self.shape.clone();
        new_shape[dim] = length;
        let out = Self::zeros(&new_shape);
        let mut idx = vec![0usize; self.shape.len()];
        let mut out_data = out.data.clone();
        loop {
            let mut src = idx.clone();
            src[dim] += start;
            let dst = idx.clone();
            out_data[out.index_at(&dst)] = self.data[self.index_at(&src)];
            if !Self::bump_indices(&mut idx, &new_shape) {
                break;
            }
        }
        Self {
            data: out_data,
            shape: new_shape,
        }
    }

    pub fn add(&self, other: &Self) -> Self {
        assert_eq!(self.data.len(), other.data.len());
        Self::from_vec(
            self.data
                .iter()
                .zip(&other.data)
                .map(|(a, b)| a + b)
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn sub(&self, other: &Self) -> Self {
        assert_eq!(self.data.len(), other.data.len());
        Self::from_vec(
            self.data
                .iter()
                .zip(&other.data)
                .map(|(a, b)| a - b)
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn mul_scalar(&self, s: f32) -> Self {
        Self::from_vec(
            self.data.iter().map(|x| x * s).collect(),
            self.shape.clone(),
        )
    }

    pub fn mul(&self, other: &Self) -> Self {
        assert_eq!(self.data.len(), other.data.len());
        Self::from_vec(
            self.data
                .iter()
                .zip(&other.data)
                .map(|(a, b)| a * b)
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn sigmoid(&self) -> Self {
        Self::from_vec(
            self.data
                .iter()
                .map(|&x| 1.0 / (1.0 + (-x).exp()))
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn gelu(&self) -> Self {
        Self::from_vec(
            self.data
                .iter()
                .map(|&x| 0.5 * x * (1.0 + erf_f32(x / std::f32::consts::SQRT_2)))
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn silu(&self) -> Self {
        Self::from_vec(
            self.data.iter().map(|&x| x / (1.0 + (-x).exp())).collect(),
            self.shape.clone(),
        )
    }

    pub fn relu(&self) -> Self {
        Self::from_vec(
            self.data.iter().map(|&x| x.max(0.0)).collect(),
            self.shape.clone(),
        )
    }

    pub fn leaky_relu(&self, leak: f32) -> Self {
        Self::from_vec(
            self.data
                .iter()
                .map(|&x| if x >= 0.0 { x } else { leak * x })
                .collect(),
            self.shape.clone(),
        )
    }

    pub fn softmax(&self, dim: usize) -> Self {
        let nd = self.shape.len();
        assert!(dim < nd);
        let outer: usize = self.shape[..dim].iter().product();
        let inner: usize = self.shape[dim + 1..].iter().product();
        let d = self.shape[dim];
        let mut out = self.data.clone();
        for o in 0..outer {
            for i in 0..inner {
                let base = o * d * inner + i;
                let mut max_v = f32::NEG_INFINITY;
                for k in 0..d {
                    max_v = max_v.max(out[base + k * inner]);
                }
                let mut sum = 0.0f32;
                for k in 0..d {
                    let idx = base + k * inner;
                    out[idx] = (out[idx] - max_v).exp();
                    sum += out[idx];
                }
                for k in 0..d {
                    let idx = base + k * inner;
                    out[idx] /= sum;
                }
            }
        }
        Self {
            data: out,
            shape: self.shape.clone(),
        }
    }

    pub fn layer_norm(&self, weight: &Tensor, bias: &Tensor, eps: f64) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, t, d) = (self.shape[0], self.shape[1], self.shape[2]);
        let mut out = vec![0.0f32; b * t * d];
        for bi in 0..b {
            for ti in 0..t {
                let base = (bi * t + ti) * d;
                let slice = &self.data[base..base + d];
                let mean: f32 = slice.iter().sum::<f32>() / d as f32;
                let var: f32 = slice.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / d as f32;
                let inv = 1.0 / (var + eps as f32).sqrt();
                for j in 0..d {
                    out[base + j] = (slice[j] - mean) * inv * weight.data[j] + bias.data[j];
                }
            }
        }
        Self {
            data: out,
            shape: self.shape.clone(),
        }
    }

    pub fn batch_norm1d(&self, weight: &Tensor, bias: &Tensor, eps: f64) -> Self {
        self.batch_norm1d_train(weight, bias, eps)
    }

    /// BatchNorm1d in eval mode using stored running statistics.
    pub fn batch_norm1d_eval(
        &self,
        weight: &Tensor,
        bias: &Tensor,
        running_mean: &Tensor,
        running_var: &Tensor,
        eps: f64,
    ) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, c, t) = (self.shape[0], self.shape[1], self.shape[2]);
        let mut out = self.data.clone();
        for ci in 0..c {
            let mean = running_mean.data[ci] as f64;
            let var = running_var.data[ci] as f64;
            let inv = 1.0 / (var + eps).sqrt();
            let w = weight.data[ci] as f64;
            let bv = bias.data[ci] as f64;
            for bi in 0..b {
                for ti in 0..t {
                    let idx = (bi * c + ci) * t + ti;
                    let v = self.data[idx] as f64;
                    out[idx] = ((v - mean) * inv * w + bv) as f32;
                }
            }
        }
        Self {
            data: out,
            shape: self.shape.clone(),
        }
    }

    pub fn batch_norm1d_train(&self, weight: &Tensor, bias: &Tensor, eps: f64) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, c, t) = (self.shape[0], self.shape[1], self.shape[2]);
        let mut out = self.data.clone();
        for ci in 0..c {
            let mut sum = 0.0f64;
            let mut sumsq = 0.0f64;
            let n = (b * t) as f64;
            for bi in 0..b {
                for ti in 0..t {
                    let v = self.data[(bi * c + ci) * t + ti] as f64;
                    sum += v;
                    sumsq += v * v;
                }
            }
            let mean = sum / n;
            let var = sumsq / n - mean * mean;
            let inv = 1.0 / (var + eps).sqrt();
            let w = weight.data[ci] as f64;
            let bv = bias.data[ci] as f64;
            for bi in 0..b {
                for ti in 0..t {
                    let idx = (bi * c + ci) * t + ti;
                    let v = self.data[idx] as f64;
                    out[idx] = ((v - mean) * inv * w + bv) as f32;
                }
            }
        }
        Self {
            data: out,
            shape: self.shape.clone(),
        }
    }

    pub fn group_norm(&self, weight: &Tensor, bias: &Tensor, num_groups: usize, eps: f64) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, c, t) = (self.shape[0], self.shape[1], self.shape[2]);
        assert_eq!(c % num_groups, 0);
        let gc = c / num_groups;
        let mut out = self.data.clone();
        for bi in 0..b {
            for gi in 0..num_groups {
                let c0 = gi * gc;
                let mut sum = 0.0f64;
                let mut sumsq = 0.0f64;
                let n = (gc * t) as f64;
                for ci in c0..c0 + gc {
                    for ti in 0..t {
                        let v = self.data[(bi * c + ci) * t + ti] as f64;
                        sum += v;
                        sumsq += v * v;
                    }
                }
                let mean = sum / n;
                let var = sumsq / n - mean * mean;
                let inv = 1.0 / (var + eps).sqrt();
                for ci in c0..c0 + gc {
                    let w = weight.data[ci] as f64;
                    let bv = bias.data[ci] as f64;
                    for ti in 0..t {
                        let idx = (bi * c + ci) * t + ti;
                        let v = self.data[idx] as f64;
                        out[idx] = ((v - mean) * inv * w + bv) as f32;
                    }
                }
            }
        }
        Self {
            data: out,
            shape: self.shape.clone(),
        }
    }

    pub fn linear(&self, weight: &Tensor, bias: Option<&Tensor>) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, t, in_d) = (self.shape[0], self.shape[1], self.shape[2]);
        let out_d = weight.shape[0];
        assert_eq!(weight.shape, vec![out_d, in_d]);
        let mut out = vec![0.0f32; b * t * out_d];
        for bi in 0..b {
            for ti in 0..t {
                let in_base = (bi * t + ti) * in_d;
                let out_base = (bi * t + ti) * out_d;
                for o in 0..out_d {
                    let mut sum = 0.0f32;
                    for i in 0..in_d {
                        sum += self.data[in_base + i] * weight.data[o * in_d + i];
                    }
                    out[out_base + o] = sum;
                }
            }
        }
        if let Some(bias) = bias {
            for bi in 0..b {
                for ti in 0..t {
                    let out_base = (bi * t + ti) * out_d;
                    for o in 0..out_d {
                        out[out_base + o] += bias.data[o];
                    }
                }
            }
        }
        Self {
            data: out,
            shape: vec![b, t, out_d],
        }
    }

    pub fn matmul_batched(&self, other: &Tensor) -> Self {
        assert_eq!(self.ndim(), 3);
        assert_eq!(other.ndim(), 3);
        let (b, m, k) = (self.shape[0], self.shape[1], self.shape[2]);
        assert_eq!(other.shape[0], b);
        let n = other.shape[2];
        assert_eq!(other.shape[1], k);
        let mut out = vec![0.0f32; b * m * n];
        for bi in 0..b {
            for i in 0..m {
                for j in 0..n {
                    let mut sum = 0.0f32;
                    for p in 0..k {
                        sum += self.data[(bi * m + i) * k + p] * other.data[(bi * k + p) * n + j];
                    }
                    out[(bi * m + i) * n + j] = sum;
                }
            }
        }
        Self {
            data: out,
            shape: vec![b, m, n],
        }
    }

    /// Conv1d on (B, C, T) layout.
    pub fn conv1d(
        &self,
        weight: &Tensor,
        bias: Option<&Tensor>,
        stride: usize,
        padding: usize,
        dilation: usize,
        groups: usize,
    ) -> Self {
        assert_eq!(self.ndim(), 3);
        let (b, c_in, t_in) = (self.shape[0], self.shape[1], self.shape[2]);
        let (c_out, c_per_g, k) = (weight.shape[0], weight.shape[1], weight.shape[2]);
        assert_eq!(c_in, c_per_g * groups);
        let t_out = (t_in + 2 * padding).saturating_sub(dilation * (k - 1) + 1) / stride + 1;
        let mut out = vec![0.0f32; b * c_out * t_out];
        for bi in 0..b {
            for co in 0..c_out {
                let g = co / (c_out / groups);
                for to in 0..t_out {
                    let mut sum = 0.0f32;
                    for ki in 0..k {
                        let ti = to * stride + ki * dilation;
                        if ti >= padding && ti < padding + t_in {
                            let t_real = ti - padding;
                            for ci in 0..c_per_g {
                                let c_in_idx = g * c_per_g + ci;
                                let w_idx = co * c_per_g * k + ci * k + ki;
                                let x_idx = (bi * c_in + c_in_idx) * t_in + t_real;
                                sum += self.data[x_idx] * weight.data[w_idx];
                            }
                        }
                    }
                    if let Some(bias) = bias {
                        sum += bias.data[co];
                    }
                    out[(bi * c_out + co) * t_out + to] = sum;
                }
            }
        }
        Self {
            data: out,
            shape: vec![b, c_out, t_out],
        }
    }

    /// Conv2d on (B, C, H, W).
    pub fn conv2d(
        &self,
        weight: &Tensor,
        bias: Option<&Tensor>,
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> Self {
        assert_eq!(self.ndim(), 4);
        let (b, c_in, h_in, w_in) = (self.shape[0], self.shape[1], self.shape[2], self.shape[3]);
        let (c_out, _, kh, kw) = (
            weight.shape[0],
            weight.shape[1],
            weight.shape[2],
            weight.shape[3],
        );
        let h_out = (h_in + 2 * padding[0] - kh) / stride[0] + 1;
        let w_out = (w_in + 2 * padding[1] - kw) / stride[1] + 1;
        let mut out = vec![0.0f32; b * c_out * h_out * w_out];
        for bi in 0..b {
            for co in 0..c_out {
                for ho in 0..h_out {
                    for wo in 0..w_out {
                        let mut sum = 0.0f32;
                        for ci in 0..c_in {
                            for ki in 0..kh {
                                for kj in 0..kw {
                                    let hi = ho * stride[0] + ki;
                                    let wi = wo * stride[1] + kj;
                                    if hi >= padding[0]
                                        && wi >= padding[1]
                                        && hi < padding[0] + h_in
                                        && wi < padding[1] + w_in
                                    {
                                        let hr = hi - padding[0];
                                        let wr = wi - padding[1];
                                        let x_idx = ((bi * c_in + ci) * h_in + hr) * w_in + wr;
                                        let w_idx = ((co * c_in + ci) * kh + ki) * kw + kj;
                                        sum += self.data[x_idx] * weight.data[w_idx];
                                    }
                                }
                            }
                        }
                        if let Some(bias) = bias {
                            sum += bias.data[co];
                        }
                        let o_idx = ((bi * c_out + co) * h_out + ho) * w_out + wo;
                        out[o_idx] = sum;
                    }
                }
            }
        }
        Self {
            data: out,
            shape: vec![b, c_out, h_out, w_out],
        }
    }

    pub fn mean_dim(&self, dim: usize, keepdim: bool) -> Self {
        let nd = self.shape.len();
        assert!(dim < nd);
        let outer: usize = self.shape[..dim].iter().product();
        let d = self.shape[dim];
        let inner: usize = self.shape[dim + 1..].iter().product();
        let mut out_shape: Vec<usize> = self.shape.clone();
        if keepdim {
            out_shape[dim] = 1;
        } else {
            out_shape.remove(dim);
        }
        let out_len: usize = out_shape.iter().product();
        let mut out = vec![0.0f32; out_len.max(1)];
        for o in 0..outer {
            for i in 0..inner {
                let mut sum = 0.0f32;
                for k in 0..d {
                    let src = if dim == 0 {
                        k * inner + i
                    } else if dim == self.shape.len() - 1 {
                        o * d + k
                    } else {
                        o * d * inner + k * inner + i
                    };
                    sum += self.data[src];
                }
                let dst = if keepdim {
                    o * inner + i
                } else if dim == 0 {
                    o * inner + i
                } else {
                    o * inner + i
                };
                if dst < out.len() {
                    out[dst] = sum / d as f32;
                }
            }
        }
        Self {
            data: out,
            shape: out_shape,
        }
    }

    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    pub fn argmax_last(&self) -> Vec<usize> {
        assert!(self.ndim() >= 1);
        let d = *self.shape.last().unwrap();
        let outer = self.data.len() / d;
        let mut out = Vec::with_capacity(outer);
        for o in 0..outer {
            let base = o * d;
            let mut best = 0usize;
            let mut best_v = self.data[base];
            for j in 1..d {
                if self.data[base + j] > best_v {
                    best_v = self.data[base + j];
                    best = j;
                }
            }
            out.push(best);
        }
        out
    }
}

/// Load tensor from tribev2-style .bin header file.
pub fn load_tensor_bin(path: &std::path::Path) -> anyhow::Result<Tensor> {
    let bytes = std::fs::read(path)?;
    let ndims = u32::from_le_bytes(bytes[0..4].try_into()?) as usize;
    let mut offset = 4;
    let mut shape = Vec::with_capacity(ndims);
    for _ in 0..ndims {
        let d = u32::from_le_bytes(bytes[offset..offset + 4].try_into()?) as usize;
        shape.push(d);
        offset += 4;
    }
    let n_floats: usize = shape.iter().product();
    let data: Vec<f32> = (0..n_floats)
        .map(|i| {
            f32::from_le_bytes(
                bytes[offset + i * 4..offset + i * 4 + 4]
                    .try_into()
                    .unwrap(),
            )
        })
        .collect();
    Ok(Tensor::from_vec(data, shape))
}

#[derive(Debug, Clone, Copy)]
pub struct TensorParityReport {
    pub pearson: f64,
    pub rmse: f64,
    pub max_abs: f32,
}

pub fn compare_tensors(a: &Tensor, b: &Tensor) -> TensorParityReport {
    assert_eq!(a.data.len(), b.data.len());
    let n = a.data.len();
    if n == 0 {
        return TensorParityReport {
            pearson: 1.0,
            rmse: 0.0,
            max_abs: 0.0,
        };
    }
    let mx: f64 = a.data.iter().map(|&v| v as f64).sum::<f64>() / n as f64;
    let my: f64 = b.data.iter().map(|&v| v as f64).sum::<f64>() / n as f64;
    let mut cov = 0.0f64;
    let mut vx = 0.0f64;
    let mut vy = 0.0f64;
    let mut se = 0.0f64;
    let mut max_abs = 0.0f32;
    for i in 0..n {
        let dx = a.data[i] as f64 - mx;
        let dy = b.data[i] as f64 - my;
        cov += dx * dy;
        vx += dx * dx;
        vy += dy * dy;
        let d = (a.data[i] - b.data[i]) as f64;
        se += d * d;
        max_abs = max_abs.max((a.data[i] - b.data[i]).abs());
    }
    let denom = (vx * vy).sqrt();
    TensorParityReport {
        pearson: if denom < 1e-30 { 1.0 } else { cov / denom },
        rmse: (se / n as f64).sqrt(),
        max_abs,
    }
}

pub fn save_tensor_bin(path: &std::path::Path, t: &Tensor) -> anyhow::Result<()> {
    let mut bytes = Vec::new();
    bytes.extend((t.shape.len() as u32).to_le_bytes());
    for &d in &t.shape {
        bytes.extend((d as u32).to_le_bytes());
    }
    for &v in &t.data {
        bytes.extend(v.to_le_bytes());
    }
    std::fs::write(path, bytes)?;
    Ok(())
}