//! WASM-compatible Tensor Operations
//!
//! CPU-only tensor operations optimized for WebAssembly execution.
//! Supports optional SIMD acceleration when available.

use serde::{Deserialize, Serialize};
use wasm_bindgen::prelude::*;

/// Data type for tensor elements
#[wasm_bindgen]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum WasmDType {
    /// 32-bit floating point
    Float32,
    /// 16-bit floating point (stored as f32)
    Float16,
    /// 8-bit integer (quantized)
    Int8,
    /// 4-bit integer (quantized, packed)
    Int4,
}

impl WasmDType {
    /// Get bytes per element
    pub fn bytes_per_element(&self) -> usize {
        match self {
            WasmDType::Float32 => 4,
            WasmDType::Float16 => 2,
            WasmDType::Int8 => 1,
            WasmDType::Int4 => 1, // Packed 2 per byte
        }
    }
}

/// A tensor for WASM computation
#[wasm_bindgen]
#[derive(Debug, Clone)]
pub struct WasmTensor {
    /// Data storage
    data: Vec<f32>,
    /// Shape
    shape: Vec<usize>,
    /// Data type
    dtype: WasmDType,
}

#[wasm_bindgen]
impl WasmTensor {
    /// Create a new tensor filled with zeros
    #[wasm_bindgen(constructor)]
    pub fn zeros(shape: Vec<usize>) -> Self {
        let size: usize = shape.iter().product();
        Self {
            data: vec![0.0; size],
            shape,
            dtype: WasmDType::Float32,
        }
    }

    /// Create a tensor from a flat array
    #[wasm_bindgen]
    pub fn from_array(data: Vec<f32>, shape: Vec<usize>) -> Result<WasmTensor, JsError> {
        let expected_size: usize = shape.iter().product();
        if data.len() != expected_size {
            return Err(JsError::new(&format!(
                "Data size {} doesn't match shape {:?} (expected {})",
                data.len(),
                shape,
                expected_size
            )));
        }
        Ok(Self {
            data,
            shape,
            dtype: WasmDType::Float32,
        })
    }

    /// Create a tensor filled with a value
    #[wasm_bindgen]
    pub fn full(shape: Vec<usize>, value: f32) -> Self {
        let size: usize = shape.iter().product();
        Self {
            data: vec![value; size],
            shape,
            dtype: WasmDType::Float32,
        }
    }

    /// Create a tensor with random values
    #[wasm_bindgen]
    pub fn randn(shape: Vec<usize>) -> Self {
        let size: usize = shape.iter().product();
        let mut data = Vec::with_capacity(size);

        // Box-Muller transform for normal distribution
        for _ in 0..size {
            let u1: f32 = rand::random();
            let u2: f32 = rand::random();
            let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
            data.push(z);
        }

        Self {
            data,
            shape,
            dtype: WasmDType::Float32,
        }
    }

    /// Get tensor shape
    #[wasm_bindgen(getter)]
    pub fn shape(&self) -> Vec<usize> {
        self.shape.clone()
    }

    /// Get number of dimensions
    #[wasm_bindgen(getter)]
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    /// Get total number of elements
    #[wasm_bindgen(getter)]
    pub fn numel(&self) -> usize {
        self.data.len()
    }

    /// Get data type
    #[wasm_bindgen(getter)]
    pub fn dtype(&self) -> WasmDType {
        self.dtype
    }

    /// Get data as array
    #[wasm_bindgen]
    pub fn to_array(&self) -> Vec<f32> {
        self.data.clone()
    }

    /// Get element at index
    #[wasm_bindgen]
    pub fn get(&self, index: usize) -> Option<f32> {
        self.data.get(index).copied()
    }

    /// Set element at index
    #[wasm_bindgen]
    pub fn set(&mut self, index: usize, value: f32) -> Result<(), JsError> {
        if index >= self.data.len() {
            return Err(JsError::new("Index out of bounds"));
        }
        self.data[index] = value;
        Ok(())
    }

    /// Reshape tensor
    #[wasm_bindgen]
    pub fn reshape(&self, new_shape: Vec<usize>) -> Result<WasmTensor, JsError> {
        let new_size: usize = new_shape.iter().product();
        if new_size != self.data.len() {
            return Err(JsError::new(&format!(
                "Cannot reshape tensor of size {} to shape {:?}",
                self.data.len(),
                new_shape
            )));
        }
        Ok(Self {
            data: self.data.clone(),
            shape: new_shape,
            dtype: self.dtype,
        })
    }

    /// Element-wise addition
    #[wasm_bindgen]
    pub fn add(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
        if self.shape != other.shape {
            return Err(JsError::new("Shape mismatch for addition"));
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(a, b)| a + b)
            .collect();
        Ok(Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        })
    }

    /// Element-wise multiplication
    #[wasm_bindgen]
    pub fn mul(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
        if self.shape != other.shape {
            return Err(JsError::new("Shape mismatch for multiplication"));
        }
        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(a, b)| a * b)
            .collect();
        Ok(Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        })
    }

    /// Scalar multiplication
    #[wasm_bindgen]
    pub fn scale(&self, scalar: f32) -> WasmTensor {
        let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
        Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// Matrix multiplication (2D tensors)
    #[wasm_bindgen]
    pub fn matmul(&self, other: &WasmTensor) -> Result<WasmTensor, JsError> {
        if self.shape.len() != 2 || other.shape.len() != 2 {
            return Err(JsError::new("matmul requires 2D tensors"));
        }
        if self.shape[1] != other.shape[0] {
            return Err(JsError::new(&format!(
                "Shape mismatch: {:?} @ {:?}",
                self.shape, other.shape
            )));
        }

        let m = self.shape[0];
        let k = self.shape[1];
        let n = other.shape[1];

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

        // Basic GEMM (can be optimized with SIMD)
        for i in 0..m {
            for j in 0..n {
                let mut sum = 0.0;
                for l in 0..k {
                    sum += self.data[i * k + l] * other.data[l * n + j];
                }
                result[i * n + j] = sum;
            }
        }

        Ok(Self {
            data: result,
            shape: vec![m, n],
            dtype: self.dtype,
        })
    }

    /// Softmax along last dimension
    #[wasm_bindgen]
    pub fn softmax(&self) -> WasmTensor {
        let last_dim = *self.shape.last().unwrap_or(&1);
        let batch_size = self.data.len() / last_dim;

        let mut result = vec![0.0; self.data.len()];

        for b in 0..batch_size {
            let start = b * last_dim;
            let end = start + last_dim;
            let slice = &self.data[start..end];

            // Find max for numerical stability
            let max_val = slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);

            // Compute exp and sum
            let exp_vals: Vec<f32> = slice.iter().map(|x| (x - max_val).exp()).collect();
            let sum: f32 = exp_vals.iter().sum();

            // Normalize
            for (i, val) in exp_vals.iter().enumerate() {
                result[start + i] = val / sum;
            }
        }

        Self {
            data: result,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// ReLU activation
    #[wasm_bindgen]
    pub fn relu(&self) -> WasmTensor {
        let data: Vec<f32> = self.data.iter().map(|x| x.max(0.0)).collect();
        Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// GELU activation
    #[wasm_bindgen]
    pub fn gelu(&self) -> WasmTensor {
        let data: Vec<f32> = self
            .data
            .iter()
            .map(|x| {
                0.5 * x
                    * (1.0
                        + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
            })
            .collect();
        Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// SiLU (Swish) activation
    #[wasm_bindgen]
    pub fn silu(&self) -> WasmTensor {
        let data: Vec<f32> = self.data.iter().map(|x| x / (1.0 + (-x).exp())).collect();
        Self {
            data,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// Layer normalization
    #[wasm_bindgen]
    pub fn layer_norm(&self, eps: f32) -> WasmTensor {
        let last_dim = *self.shape.last().unwrap_or(&1);
        let batch_size = self.data.len() / last_dim;

        let mut result = vec![0.0; self.data.len()];

        for b in 0..batch_size {
            let start = b * last_dim;
            let end = start + last_dim;
            let slice = &self.data[start..end];

            // Compute mean
            let mean: f32 = slice.iter().sum::<f32>() / last_dim as f32;

            // Compute variance
            let var: f32 = slice.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / last_dim as f32;

            // Normalize
            let std = (var + eps).sqrt();
            for (i, val) in slice.iter().enumerate() {
                result[start + i] = (val - mean) / std;
            }
        }

        Self {
            data: result,
            shape: self.shape.clone(),
            dtype: self.dtype,
        }
    }

    /// Sum all elements
    #[wasm_bindgen]
    pub fn sum(&self) -> f32 {
        self.data.iter().sum()
    }

    /// Mean of all elements
    #[wasm_bindgen]
    pub fn mean(&self) -> f32 {
        self.sum() / self.data.len() as f32
    }

    /// Max of all elements
    #[wasm_bindgen]
    pub fn max(&self) -> f32 {
        self.data.iter().copied().fold(f32::NEG_INFINITY, f32::max)
    }

    /// Argmax (index of maximum element)
    #[wasm_bindgen]
    pub fn argmax(&self) -> usize {
        self.data
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map_or(0, |(i, _)| i)
    }

    /// Memory usage in bytes
    #[wasm_bindgen]
    pub fn memory_usage(&self) -> usize {
        self.data.len() * std::mem::size_of::<f32>()
    }
}

/// KV-Cache for transformer inference
#[wasm_bindgen]
#[derive(Debug, Clone)]
pub struct WasmKvCache {
    /// Key cache: [num_layers, num_heads, max_seq_len, head_dim]
    keys: Vec<f32>,
    /// Value cache: [num_layers, num_heads, max_seq_len, head_dim]
    values: Vec<f32>,
    /// Configuration
    max_seq_len: usize,
    num_layers: usize,
    num_heads: usize,
    head_dim: usize,
    /// Current position
    position: usize,
}

#[wasm_bindgen]
impl WasmKvCache {
    /// Create a new KV cache
    #[wasm_bindgen(constructor)]
    pub fn new(max_seq_len: usize, num_layers: usize, num_heads: usize, head_dim: usize) -> Self {
        let cache_size = num_layers * num_heads * max_seq_len * head_dim;
        Self {
            keys: vec![0.0; cache_size],
            values: vec![0.0; cache_size],
            max_seq_len,
            num_layers,
            num_heads,
            head_dim,
            position: 0,
        }
    }

    /// Get current position
    #[wasm_bindgen(getter)]
    pub fn position(&self) -> usize {
        self.position
    }

    /// Advance position by one
    pub fn step(&mut self) {
        self.position = (self.position + 1).min(self.max_seq_len - 1);
    }

    /// Clear the cache
    #[wasm_bindgen]
    pub fn clear(&mut self) {
        self.keys.fill(0.0);
        self.values.fill(0.0);
        self.position = 0;
    }

    /// Memory usage in bytes
    #[wasm_bindgen]
    pub fn memory_usage(&self) -> usize {
        (self.keys.len() + self.values.len()) * std::mem::size_of::<f32>()
    }

    /// Get cache utilization (0.0 - 1.0)
    #[wasm_bindgen]
    pub fn utilization(&self) -> f32 {
        self.position as f32 / self.max_seq_len as f32
    }
}

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

    #[test]
    fn test_tensor_zeros() {
        let t = WasmTensor::zeros(vec![2, 3]);
        assert_eq!(t.shape(), vec![2, 3]);
        assert_eq!(t.numel(), 6);
        assert_eq!(t.sum(), 0.0);
    }

    #[test]
    fn test_tensor_from_array() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let t = WasmTensor::from_array(data, vec![2, 3]).unwrap();
        assert_eq!(t.get(0), Some(1.0));
        assert_eq!(t.get(5), Some(6.0));
    }

    #[test]
    fn test_tensor_matmul() {
        // 2x3 @ 3x2 = 2x2
        let a = WasmTensor::from_array(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let b = WasmTensor::from_array(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.shape(), vec![2, 2]);
        // [1,2,3] @ [1,2] = 1+6+15 = 22
        assert!((c.get(0).unwrap() - 22.0).abs() < 1e-5);
    }

    #[test]
    fn test_tensor_softmax() {
        let t = WasmTensor::from_array(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
        let s = t.softmax();
        assert!((s.sum() - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_kv_cache() {
        let mut cache = WasmKvCache::new(1024, 32, 32, 128);
        assert_eq!(cache.position(), 0);
        cache.step();
        assert_eq!(cache.position(), 1);
        cache.clear();
        assert_eq!(cache.position(), 0);
    }
}