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

pub mod activations;
pub use activations::{relu, gelu, softmax};

pub fn add(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.shape() != b.shape() {
        return Err("Shapes must match for addition".to_string());
    }
    if a.dtype() != b.dtype() {
        return Err("DTypes must match".to_string());
    }
    if a.device() != DeviceType::Cpu || b.device() != DeviceType::Cpu {
        return Err("Only CPU add supported for now".to_string());
    }

    let count = a.shape().num_elements();
    let mut result = Tensor::new(a.shape().clone(), a.dtype(), a.device())?;

    match a.dtype() {
        DType::F32 => {
            let a_data: Vec<f32> = a.read_data()?;
            let b_data: Vec<f32> = b.read_data()?;
            
            // Parallel addition using Rayon
            let mut res_data = vec![0.0; count];
            res_data.par_iter_mut().enumerate().for_each(|(i, x)| {
                *x = a_data[i] + b_data[i];
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType for add".to_string()),
    }

    Ok(result)
}

pub fn sub(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.shape() != b.shape() { return Err("Shapes must match".to_string()); }
    if a.dtype() != b.dtype() { return Err("DTypes must match".to_string()); }
    
    let count = a.shape().num_elements();
    let mut result = Tensor::new(a.shape().clone(), a.dtype(), a.device())?;

    match a.dtype() {
        DType::F32 => {
            let a_data: Vec<f32> = a.read_data()?;
            let b_data: Vec<f32> = b.read_data()?;
            
            let mut res_data = vec![0.0; count];
            res_data.par_iter_mut().enumerate().for_each(|(i, x)| {
                *x = a_data[i] - b_data[i];
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType".to_string()),
    }
    Ok(result)
}

pub fn mul(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.shape() != b.shape() { return Err("Shapes must match".to_string()); }
    if a.dtype() != b.dtype() { return Err("DTypes must match".to_string()); }
    
    let count = a.shape().num_elements();
    let mut result = Tensor::new(a.shape().clone(), a.dtype(), a.device())?;

    match a.dtype() {
        DType::F32 => {
            let a_data: Vec<f32> = a.read_data()?;
            let b_data: Vec<f32> = b.read_data()?;
            
            let mut res_data = vec![0.0; count];
            res_data.par_iter_mut().enumerate().for_each(|(i, x)| {
                *x = a_data[i] * b_data[i];
            });

            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType".to_string()),
    }
    Ok(result)
}

pub fn div(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.shape() != b.shape() { return Err("Shapes must match".to_string()); }
    if a.dtype() != b.dtype() { return Err("DTypes must match".to_string()); }
    
    let count = a.shape().num_elements();
    let mut result = Tensor::new(a.shape().clone(), a.dtype(), a.device())?;

    match a.dtype() {
        DType::F32 => {
            let a_data: Vec<f32> = a.read_data()?;
            let b_data: Vec<f32> = b.read_data()?;
            
            let mut res_data = vec![0.0; count];
            res_data.par_iter_mut().enumerate().for_each(|(i, x)| {
                *x = a_data[i] / b_data[i];
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType".to_string()),
    }
    Ok(result)
}

pub fn matmul(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.dtype() != b.dtype() { return Err("DTypes must match".to_string()); }
    
    // Check shapes: (M, K) @ (K, N) -> (M, N)
    if a.shape().rank() != 2 || b.shape().rank() != 2 {
        return Err("Matmul currently supports only rank-2 tensors".to_string());
    }
    
    let m = a.shape().dims()[0];
    let k = a.shape().dims()[1];
    let k2 = b.shape().dims()[0];
    let n = b.shape().dims()[1];

    if k != k2 {
        return Err(format!("Shape mismatch for matmul: ({}, {}) vs ({}, {})", m, k, k2, n));
    }

    let out_shape = crate::tensor::Shape::new(vec![m, n]);
    let mut result = Tensor::new(out_shape, a.dtype(), a.device())?;

    match a.dtype() {
        DType::F32 => {
            let a_data: Vec<f32> = a.read_data()?;
            let b_data: Vec<f32> = b.read_data()?;
            let mut res_data = vec![0.0; m * n];

            // Parallelized MatMul
            // Parallelize over the rows of the output matrix (m)
            res_data.chunks_mut(n).enumerate().for_each(|(i, row)| {
                for j in 0..n {
                    let mut sum = 0.0;
                    for p in 0..k {
                        sum += a_data[i * k + p] * b_data[p * n + j];
                    }
                    row[j] = sum;
                }
            });
            
            result.write_data(&res_data)?;
        },
        _ => return Err("Unsupported DType".to_string()),
    }
    Ok(result)
}