use numpy::{PyArray1, PyReadonlyArray1, PyReadonlyArray2};
use pyo3::prelude::*;

use crate::utils::sigmoid;

/// Logistic regression batch prediction.
///
/// Given a feature matrix (N x F), weight vector (F,), and bias scalar,
/// computes sigmoid(X @ weights + bias) for each row.
/// Returns a numpy array of probabilities.
///
/// Uses ndarray's vectorized matrix-vector multiply (GEMV) which
/// auto-vectorizes with SIMD, instead of per-row dot products.
#[pyfunction]
pub fn predict_batch<'py>(
    py: Python<'py>,
    features: PyReadonlyArray2<f64>,
    weights: PyReadonlyArray1<f64>,
    bias: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let array = features.as_array();
    let w = weights.as_array();
    let n_cols = array.ncols();

    if n_cols != w.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "Feature matrix has {} columns but weights has {} elements",
            n_cols,
            w.len()
        )));
    }

    let probs: Vec<f64> = py.allow_threads(|| {
        // Vectorized matrix-vector multiply: logits = X @ w + bias
        let logits = array.dot(&w) + bias;
        // Apply sigmoid element-wise
        logits.mapv(|x| sigmoid(x)).to_vec()
    });

    Ok(PyArray1::from_vec(py, probs))
}

#[cfg(test)]
mod tests {
    use crate::utils::sigmoid;

    #[test]
    fn test_dot_product_sigmoid() {
        let features = vec![1.0, 2.0, 3.0];
        let weights = vec![0.5, 0.5, 0.5];
        let bias = -1.0;
        let dot: f64 = features.iter().zip(weights.iter()).map(|(x, w)| x * w).sum();
        let prob = sigmoid(dot + bias);
        assert!((prob - 0.8808).abs() < 0.001);
    }
}
