//! Comparison operations for language-compatible logic
//!
//! Implements comparison operators returning logical matrices/values.

use runmat_builtins::Tensor;

/// Element-wise greater than comparison
pub fn matrix_gt(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} > {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| if x > y { 1.0 } else { 0.0 })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}

/// Element-wise greater than or equal comparison
pub fn matrix_ge(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} >= {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| if x >= y { 1.0 } else { 0.0 })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}

/// Element-wise less than comparison
pub fn matrix_lt(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} < {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| if x < y { 1.0 } else { 0.0 })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}

/// Element-wise less than or equal comparison
pub fn matrix_le(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} <= {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| if x <= y { 1.0 } else { 0.0 })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}

/// Element-wise equality comparison
pub fn matrix_eq(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} == {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| {
            if (x - y).abs() < f64::EPSILON {
                1.0
            } else {
                0.0
            }
        })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}

/// Element-wise inequality comparison
pub fn matrix_ne(a: &Tensor, b: &Tensor) -> Result<Tensor, String> {
    if a.rows() != b.rows() || a.cols() != b.cols() {
        return Err(format!(
            "Matrix dimensions must agree: {}x{} != {}x{}",
            a.rows(),
            a.cols(),
            b.rows(),
            b.cols()
        ));
    }

    let data: Vec<f64> = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| {
            if (x - y).abs() >= f64::EPSILON {
                1.0
            } else {
                0.0
            }
        })
        .collect();

    Tensor::new_2d(data, a.rows(), a.cols())
}
