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

/// SVM batch prediction using RBF kernel with Platt scaling.
///
/// For each sample x in the feature matrix (N x F), computes:
///   decision = intercept + sum_i(dual_coef[i] * exp(-gamma * ||x - sv[i]||^2))
///   probability = 1 / (1 + exp(prob_a * decision + prob_b))
///
/// The prob_a and prob_b parameters come from sklearn's Platt calibration
/// (SVC.probA_[0] and SVC.probB_[0]).
#[pyfunction]
pub fn predict_svm_batch<'py>(
    py: Python<'py>,
    features: PyReadonlyArray2<f64>,
    support_vectors: PyReadonlyArray2<f64>,
    dual_coef: PyReadonlyArray1<f64>,
    intercept: f64,
    gamma: f64,
    prob_a: f64,
    prob_b: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let feat = features.as_array();
    let sv = support_vectors.as_array();
    let coef = dual_coef.as_array();

    let n_rows = feat.nrows();
    let n_sv = sv.nrows();
    let n_feat = feat.ncols();

    if sv.ncols() != n_feat {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "Feature matrix has {} columns but support vectors have {} columns",
            n_feat,
            sv.ncols()
        )));
    }
    if coef.len() != n_sv {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "support_vectors has {} rows but dual_coef has {} elements",
            n_sv,
            coef.len()
        )));
    }

    // Ensure contiguous layout for slicing
    let owned_feat;
    let feat_ref = if feat.as_slice().is_some() {
        feat
    } else {
        owned_feat = feat.as_standard_layout().into_owned();
        owned_feat.view()
    };
    let owned_sv;
    let sv_ref = if sv.as_slice().is_some() {
        sv
    } else {
        owned_sv = sv.as_standard_layout().into_owned();
        owned_sv.view()
    };

    let rows: Vec<&[f64]> = (0..n_rows)
        .map(|i| feat_ref.row(i).to_slice().unwrap())
        .collect();
    let sv_rows: Vec<&[f64]> = (0..n_sv)
        .map(|i| sv_ref.row(i).to_slice().unwrap())
        .collect();
    let coef_slice = coef.as_slice().unwrap();

    // Precompute squared norms of support vectors to avoid redundant work.
    // ||x - sv||^2 = ||x||^2 + ||sv||^2 - 2*(x . sv)
    let sv_norms_sq: Vec<f64> = sv_rows
        .iter()
        .map(|sv| sv.iter().map(|v| v * v).sum())
        .collect();

    let probs: Vec<f64> = py.allow_threads(|| {
        rows.par_iter()
            .map(|sample| {
                let x_norm_sq: f64 = sample.iter().map(|v| v * v).sum();
                let mut decision = intercept;
                for (i, sv_row) in sv_rows.iter().enumerate() {
                    let dot: f64 = sample
                        .iter()
                        .zip(sv_row.iter())
                        .map(|(a, b)| a * b)
                        .sum();
                    let dist_sq = x_norm_sq + sv_norms_sq[i] - 2.0 * dot;
                    decision += coef_slice[i] * (-gamma * dist_sq).exp();
                }
                // Platt scaling: P(y=1|f) = 1 / (1 + exp(A*f - B))
                // Note: B is negated to match sklearn/libsvm's internal class ordering
                1.0 / (1.0 + (prob_a * decision - prob_b).exp())
            })
            .collect()
    });

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

#[cfg(test)]
mod tests {
    /// Platt scaling matching sklearn/libsvm: P(y=1|f) = 1 / (1 + exp(A*f - B))
    fn platt(decision: f64, prob_a: f64, prob_b: f64) -> f64 {
        1.0 / (1.0 + (prob_a * decision - prob_b).exp())
    }

    #[test]
    fn test_svm_rbf_single_sample() {
        // Single sample [1.0, 0.0], single support vector [0.0, 0.0]
        // gamma=1.0, dual_coef=1.0, intercept=0.0
        // K(x, sv) = exp(-1.0 * 1.0) = exp(-1.0) ≈ 0.3679
        // decision = 0.0 + 1.0 * 0.3679 = 0.3679
        // With prob_a=-1.0, prob_b=0.0 (equivalent to plain sigmoid):
        //   prob = 1/(1+exp(-0.3679)) ≈ 0.5909
        let gamma = 1.0;
        let dist_sq: f64 = 1.0;
        let kernel_val = (-gamma * dist_sq).exp();
        let decision = 0.0 + 1.0 * kernel_val;
        let prob = platt(decision, -1.0, 0.0);
        assert!((kernel_val - 0.3679).abs() < 0.001);
        assert!((prob - 0.5909).abs() < 0.001);
    }

    #[test]
    fn test_svm_rbf_two_support_vectors() {
        // Sample [0.5, 0.5], two SVs: [0,0] and [1,1]
        // gamma=0.5, dual_coef=[1.0, -1.0], intercept=0.0
        // decision = 0.0 (symmetric)
        // With prob_a=-1.0, prob_b=0.0: prob = 0.5
        let gamma = 0.5;
        let d0: f64 = (0.5_f64).powi(2) + (0.5_f64).powi(2);
        let d1: f64 = (0.5 - 1.0_f64).powi(2) + (0.5 - 1.0_f64).powi(2);
        let k0 = (-gamma * d0).exp();
        let k1 = (-gamma * d1).exp();
        let decision = 1.0 * k0 + (-1.0) * k1;
        let prob = platt(decision, -1.0, 0.0);
        assert!((decision).abs() < 1e-10);
        assert!((prob - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_platt_scaling() {
        // Typical sklearn Platt params: prob_a ≈ -5.0, prob_b ≈ 0.0
        // decision=2.0 → prob = 1/(1+exp(-5*2+0)) = 1/(1+exp(-10)) ≈ 1.0
        // decision=-2.0 → prob = 1/(1+exp(10)) ≈ 0.0
        let p_high = platt(2.0, -5.0, 0.0);
        let p_low = platt(-2.0, -5.0, 0.0);
        assert!(p_high > 0.999);
        assert!(p_low < 0.001);
    }
}
