//! CPU tensor operations for LLMHosts Core inference.

use thiserror::Error;

#[derive(Error, Debug)]
pub enum TensorError {
    #[error("Shape mismatch: expected {expected:?}, got {got:?}")]
    ShapeMismatch { expected: Vec<usize>, got: Vec<usize> },
    #[error("Empty tensor")]
    Empty,
}

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

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

    pub fn from_data(data: Vec<f32>, shape: Vec<usize>) -> Result<Self, TensorError> {
        let expected: usize = shape.iter().product();
        if data.len() != expected {
            return Err(TensorError::ShapeMismatch {
                expected: shape,
                got: vec![data.len()],
            });
        }
        Ok(Self { data, shape })
    }

    /// Element-wise add (broadcast 1D weight over rows if shapes differ).
    pub fn add_inplace(&mut self, other: &Tensor) -> Result<(), TensorError> {
        if self.data.len() == other.data.len() {
            // Same total size — element-wise
            for (a, b) in self.data.iter_mut().zip(other.data.iter()) {
                *a += b;
            }
        } else if self.shape.len() == 2 && other.shape.len() == 1 && self.shape[1] == other.shape[0] {
            // Broadcast: [seq_len, hidden] + [hidden]
            let cols = other.shape[0];
            for row in self.data.chunks_mut(cols) {
                for (a, b) in row.iter_mut().zip(other.data.iter()) {
                    *a += b;
                }
            }
        } else {
            return Err(TensorError::ShapeMismatch {
                expected: self.shape.clone(),
                got: other.shape.clone(),
            });
        }
        Ok(())
    }

    /// Element-wise multiply.
    pub fn mul_inplace(&mut self, other: &Tensor) -> Result<(), TensorError> {
        if self.data.len() != other.data.len() {
            return Err(TensorError::ShapeMismatch {
                expected: self.shape.clone(),
                got: other.shape.clone(),
            });
        }
        for (a, b) in self.data.iter_mut().zip(other.data.iter()) {
            *a *= b;
        }
        Ok(())
    }

    /// Scale all elements.
    pub fn scale_inplace(&mut self, factor: f32) {
        for v in &mut self.data {
            *v *= factor;
        }
    }

    /// Softmax over the last dimension.
    pub fn softmax_last_dim(&mut self) {
        if self.shape.is_empty() {
            return;
        }
        let last_dim = *self.shape.last().unwrap();
        for chunk in self.data.chunks_mut(last_dim) {
            let max_val = chunk.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let mut sum = 0.0f32;
            for v in chunk.iter_mut() {
                *v = (*v - max_val).exp();
                sum += *v;
            }
            if sum > 0.0 {
                for v in chunk.iter_mut() {
                    *v /= sum;
                }
            }
        }
    }

    /// RMS normalization (returns new tensor, does NOT apply learned weight).
    pub fn rms_norm(&self, eps: f32) -> Tensor {
        if self.shape.is_empty() {
            return self.clone();
        }
        let hidden = *self.shape.last().unwrap();
        let mut out = self.data.clone();
        for chunk in out.chunks_mut(hidden) {
            let ss: f32 = chunk.iter().map(|x| x * x).sum::<f32>() / hidden as f32;
            let rms = (ss + eps).sqrt();
            for v in chunk.iter_mut() {
                *v /= rms;
            }
        }
        Tensor {
            data: out,
            shape: self.shape.clone(),
        }
    }
}

/// SiLU activation in-place: x * sigmoid(x).
pub fn silu_inplace(t: &mut Tensor) {
    for v in &mut t.data {
        *v = *v * (1.0 / (1.0 + (-*v).exp()));
    }
}

/// Matrix multiply two 2D tensors: A [M, K] x B [K, N] → C [M, N].
///
/// Uses matrixmultiply::sgemm for 10-20x speedup over naive triple-loop.
pub fn matmul_2d(a: &Tensor, b: &Tensor) -> Result<Tensor, TensorError> {
    if a.shape.len() != 2 || b.shape.len() != 2 {
        return Err(TensorError::ShapeMismatch {
            expected: vec![2],
            got: a.shape.clone(),
        });
    }
    let (m, k) = (a.shape[0], a.shape[1]);
    let (k2, n) = (b.shape[0], b.shape[1]);
    if k != k2 {
        return Err(TensorError::ShapeMismatch {
            expected: vec![m, k],
            got: b.shape.clone(),
        });
    }

    let mut c = vec![0.0f32; m * n];

    // sgemm: C = alpha * A * B + beta * C
    // Row-major: rsa=K, csa=1 for A; rsb=N, csb=1 for B; rsc=N, csc=1 for C
    unsafe {
        matrixmultiply::sgemm(
            m,
            k,
            n,
            1.0,
            a.data.as_ptr(),
            k as isize,
            1,
            b.data.as_ptr(),
            n as isize,
            1,
            0.0,
            c.as_mut_ptr(),
            n as isize,
            1,
        );
    }

    Ok(Tensor {
        data: c,
        shape: vec![m, n],
    })
}

/// Gather rows from a 2D tensor by index. Used for embedding lookup.
/// input shape: [vocab_size, hidden_size], indices: list of token IDs
/// output shape: [seq_len, hidden_size]
pub fn gather_rows(weight: &Tensor, indices: &[u32]) -> Result<Tensor, TensorError> {
    if weight.shape.len() != 2 {
        return Err(TensorError::ShapeMismatch {
            expected: vec![2],
            got: weight.shape.clone(),
        });
    }
    let vocab = weight.shape[0];
    let hidden = weight.shape[1];
    let mut out = Vec::with_capacity(indices.len() * hidden);
    for &idx in indices {
        let row = idx as usize;
        if row >= vocab {
            return Err(TensorError::ShapeMismatch {
                expected: vec![vocab],
                got: vec![row],
            });
        }
        let start = row * hidden;
        out.extend_from_slice(&weight.data[start..start + hidden]);
    }
    Ok(Tensor {
        data: out,
        shape: vec![indices.len(), hidden],
    })
}

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

    #[test]
    fn test_zeros() {
        let t = Tensor::zeros(vec![2, 3]);
        assert_eq!(t.data.len(), 6);
        assert!(t.data.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn test_from_data_valid() {
        let t = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        assert_eq!(t.shape, vec![2, 2]);
    }

    #[test]
    fn test_from_data_invalid() {
        let result = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![2, 2]);
        assert!(result.is_err());
    }

    #[test]
    fn test_add_inplace_same_shape() {
        let mut a = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
        let b = Tensor::from_data(vec![10.0, 20.0, 30.0], vec![3]).unwrap();
        a.add_inplace(&b).unwrap();
        assert_eq!(a.data, vec![11.0, 22.0, 33.0]);
    }

    #[test]
    fn test_add_inplace_broadcast_1d() {
        let mut a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
        let b = Tensor::from_data(vec![10.0, 20.0], vec![2]).unwrap();
        a.add_inplace(&b).unwrap();
        assert_eq!(a.data, vec![11.0, 22.0, 13.0, 24.0]);
    }

    #[test]
    fn test_mul_inplace() {
        let mut a = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
        let b = Tensor::from_data(vec![2.0, 3.0, 4.0], vec![3]).unwrap();
        a.mul_inplace(&b).unwrap();
        assert_eq!(a.data, vec![2.0, 6.0, 12.0]);
    }

    #[test]
    fn test_softmax_last_dim() {
        let mut t = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
        t.softmax_last_dim();
        let sum: f32 = t.data.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
        // Values should be increasing
        assert!(t.data[0] < t.data[1]);
        assert!(t.data[1] < t.data[2]);
    }

    #[test]
    fn test_rms_norm() {
        let t = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        let normed = t.rms_norm(1e-6);
        let ss: f32 = normed.data.iter().map(|x| x * x).sum::<f32>() / 4.0;
        // After RMS norm, the mean of squares should be ~1.0
        assert!((ss - 1.0).abs() < 0.01, "ss = {ss}");
    }

    #[test]
    fn test_silu_inplace() {
        let mut t = Tensor::from_data(vec![0.0, 1.0, -1.0], vec![3]).unwrap();
        silu_inplace(&mut t);
        // silu(0) = 0, silu(1) ≈ 0.731, silu(-1) ≈ -0.269
        assert!((t.data[0] - 0.0).abs() < 0.01);
        assert!((t.data[1] - 0.731).abs() < 0.01);
        assert!((t.data[2] - (-0.269)).abs() < 0.01);
    }

    #[test]
    fn test_matmul_2d() {
        // [2,3] x [3,2] = [2,2]
        let a = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let b = Tensor::from_data(vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0], vec![3, 2]).unwrap();
        let c = matmul_2d(&a, &b).unwrap();
        assert_eq!(c.shape, vec![2, 2]);
        // Row 0: 1*7+2*9+3*11=58, 1*8+2*10+3*12=64
        assert!((c.data[0] - 58.0).abs() < 0.01);
        assert!((c.data[1] - 64.0).abs() < 0.01);
        // Row 1: 4*7+5*9+6*11=139, 4*8+5*10+6*12=154
        assert!((c.data[2] - 139.0).abs() < 0.01);
        assert!((c.data[3] - 154.0).abs() < 0.01);
    }

    #[test]
    fn test_matmul_shape_mismatch() {
        let a = Tensor::from_data(vec![1.0; 6], vec![2, 3]).unwrap();
        let b = Tensor::from_data(vec![1.0; 8], vec![4, 2]).unwrap();
        assert!(matmul_2d(&a, &b).is_err());
    }

    #[test]
    fn test_gather_rows() {
        let w = Tensor::from_data(
            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0],
            vec![4, 3],
        ).unwrap();
        let result = gather_rows(&w, &[0, 2]).unwrap();
        assert_eq!(result.shape, vec![2, 3]);
        assert_eq!(result.data, vec![1.0, 2.0, 3.0, 7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_scale_inplace() {
        let mut t = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
        t.scale_inplace(2.0);
        assert_eq!(t.data, vec![2.0, 4.0, 6.0]);
    }
}
