use crate::tensor::{Tensor, DType};
use rayon::prelude::*;

pub fn relu(t: &Tensor) -> Result<Tensor, String> {
    let mut result = Tensor::new(t.shape().clone(), t.dtype(), t.device())?;
    
    match t.dtype() {
        DType::F32 => {
            let data: Vec<f32> = t.read_data()?;
            let mut res_data = vec![0.0; data.len()];
            
            res_data.par_iter_mut().enumerate().for_each(|(i, x)| {
                *x = if data[i] > 0.0 { data[i] } else { 0.0 };
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType for ReLU".to_string()),
    }
    Ok(result)
}

pub fn gelu(t: &Tensor) -> Result<Tensor, String> {
    let mut result = Tensor::new(t.shape().clone(), t.dtype(), t.device())?;
    
    match t.dtype() {
        DType::F32 => {
            let data: Vec<f32> = t.read_data()?;
            let mut res_data = vec![0.0; data.len()];
            
            // Approximation of GELU: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
            res_data.par_iter_mut().enumerate().for_each(|(i, val)| {
                let x = data[i];
                let c = 0.7978845608; // sqrt(2/pi)
                let tanh_inner = c * (x + 0.044715 * x * x * x);
                *val = 0.5 * x * (1.0 + tanh_inner.tanh());
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType for GELU".to_string()),
    }
    Ok(result)
}

pub fn softmax(t: &Tensor, axis: usize) -> Result<Tensor, String> {
    // Basic Softmax implementation (Currently treats as flattened if axis logic is complex)
    // For now, let's implement a safe softmax for the last dimension (common in LLMs)
    
    if t.shape().rank() != 2 {
        return Err("Softmax currently only optimized for rank-2 (Batch, Vocab)".to_string());
    }
    if axis != 1 {
        return Err("Softmax currently only supports axis 1".to_string());
    }

    let mut result = Tensor::new(t.shape().clone(), t.dtype(), t.device())?;
    let _rows = t.shape().dims()[0];
    let cols = t.shape().dims()[1];

    match t.dtype() {
        DType::F32 => {
            let data: Vec<f32> = t.read_data()?;
            let mut res_data = vec![0.0; data.len()];
            
            res_data.chunks_mut(cols).enumerate().for_each(|(i, row_out)| {
                let row_in = &data[i * cols .. (i + 1) * cols];
                
                // Numerical stability: subtract max
                let max_val = row_in.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
                let mut sum_exp = 0.0;
                
                for (j, val) in row_in.iter().enumerate() {
                    let exp_val = (val - max_val).exp();
                    row_out[j] = exp_val;
                    sum_exp += exp_val;
                }
                
                for val in row_out.iter_mut() {
                    *val /= sum_exp;
                }
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType for Softmax".to_string()),
    }
    Ok(result)
}
