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

use crate::utils::sigmoid;

/// Traverse a single decision tree for one sample.
///
/// Returns the leaf value (class-1 probability for RF, raw prediction for GBDT).
fn traverse_tree(
    sample: &[f64],
    children_left: &[i64],
    children_right: &[i64],
    features: &[i64],
    thresholds: &[f64],
    values: &[f64],
) -> f64 {
    let mut node: usize = 0;
    loop {
        let left = children_left[node];
        // Leaf node: children_left == -1 (sklearn convention)
        if left < 0 {
            return values[node];
        }
        let feat_idx = features[node] as usize;
        let threshold = thresholds[node];
        let feat_val = if feat_idx < sample.len() {
            sample[feat_idx]
        } else {
            f64::NAN
        };
        // sklearn convention: go left if <= threshold or NaN
        if feat_val <= threshold || feat_val.is_nan() {
            node = left as usize;
        } else {
            node = children_right[node] as usize;
        }
    }
}

/// Random Forest batch prediction.
///
/// Tree structures are passed as flat numpy arrays with an offsets array.
/// offsets[t]..offsets[t+1] delimits tree t in each flat array.
/// This avoids repeated Python list → Rust Vec conversion on every predict call.
#[pyfunction]
pub fn predict_forest_batch<'py>(
    py: Python<'py>,
    features: PyReadonlyArray2<f64>,
    offsets: PyReadonlyArray1<i64>,
    children_left: PyReadonlyArray1<i64>,
    children_right: PyReadonlyArray1<i64>,
    tree_features: PyReadonlyArray1<i64>,
    thresholds: PyReadonlyArray1<f64>,
    values: PyReadonlyArray1<f64>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let array = features.as_array();
    let n_rows = array.nrows();
    let off = offsets.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("offsets must be contiguous: {e}"))
    })?;
    let cl = children_left.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("children_left must be contiguous: {e}"))
    })?;
    let cr = children_right.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("children_right must be contiguous: {e}"))
    })?;
    let tf = tree_features.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("tree_features must be contiguous: {e}"))
    })?;
    let th = thresholds.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("thresholds must be contiguous: {e}"))
    })?;
    let val = values.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("values must be contiguous: {e}"))
    })?;

    let n_trees = off.len().saturating_sub(1);
    if n_trees == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "Must provide at least one tree",
        ));
    }

    // Ensure contiguous layout — to_slice() returns None for non-contiguous arrays
    let owned;
    let array_ref = if array.as_slice().is_some() {
        array
    } else {
        owned = array.as_standard_layout().into_owned();
        owned.view()
    };
    let rows: Vec<&[f64]> = (0..n_rows)
        .map(|i| array_ref.row(i).to_slice().unwrap())
        .collect();

    let probs: Vec<f64> = py.allow_threads(|| {
        rows.par_iter()
            .map(|sample| {
                let mut sum = 0.0;
                for t in 0..n_trees {
                    let start = off[t] as usize;
                    let end = off[t + 1] as usize;
                    sum += traverse_tree(
                        sample,
                        &cl[start..end],
                        &cr[start..end],
                        &tf[start..end],
                        &th[start..end],
                        &val[start..end],
                    );
                }
                sum / n_trees as f64
            })
            .collect()
    });

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

/// Gradient Boosting batch prediction.
///
/// Tree structures are passed as flat numpy arrays with an offsets array.
/// For each sample, sums learning_rate * leaf_value across all trees,
/// adds init_value, and applies sigmoid.
#[pyfunction]
pub fn predict_gbdt_batch<'py>(
    py: Python<'py>,
    features: PyReadonlyArray2<f64>,
    offsets: PyReadonlyArray1<i64>,
    children_left: PyReadonlyArray1<i64>,
    children_right: PyReadonlyArray1<i64>,
    tree_features: PyReadonlyArray1<i64>,
    thresholds: PyReadonlyArray1<f64>,
    values: PyReadonlyArray1<f64>,
    learning_rate: f64,
    init_value: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let array = features.as_array();
    let n_rows = array.nrows();
    let off = offsets.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("offsets must be contiguous: {e}"))
    })?;
    let cl = children_left.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("children_left must be contiguous: {e}"))
    })?;
    let cr = children_right.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("children_right must be contiguous: {e}"))
    })?;
    let tf = tree_features.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("tree_features must be contiguous: {e}"))
    })?;
    let th = thresholds.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("thresholds must be contiguous: {e}"))
    })?;
    let val = values.as_slice().map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("values must be contiguous: {e}"))
    })?;

    let n_trees = off.len().saturating_sub(1);
    if n_trees == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "Must provide at least one tree",
        ));
    }

    let owned_gbdt;
    let array_ref_gbdt = if array.as_slice().is_some() {
        array
    } else {
        owned_gbdt = array.as_standard_layout().into_owned();
        owned_gbdt.view()
    };
    let rows: Vec<&[f64]> = (0..n_rows)
        .map(|i| array_ref_gbdt.row(i).to_slice().unwrap())
        .collect();

    let probs: Vec<f64> = py.allow_threads(|| {
        rows.par_iter()
            .map(|sample| {
                let mut raw = init_value;
                for t in 0..n_trees {
                    let start = off[t] as usize;
                    let end = off[t + 1] as usize;
                    raw += learning_rate
                        * traverse_tree(
                            sample,
                            &cl[start..end],
                            &cr[start..end],
                            &tf[start..end],
                            &th[start..end],
                            &val[start..end],
                        );
                }
                sigmoid(raw)
            })
            .collect()
    });

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_traverse_single_leaf() {
        // A tree with just one leaf (root is leaf)
        let children_left = vec![-1];
        let children_right = vec![-1];
        let features = vec![-2]; // undefined for leaf
        let thresholds = vec![-2.0];
        let values = vec![0.75];

        let sample = vec![1.0, 2.0];
        let result = traverse_tree(&sample, &children_left, &children_right, &features, &thresholds, &values);
        assert!((result - 0.75).abs() < 1e-10);
    }

    #[test]
    fn test_traverse_simple_split() {
        // Root splits on feature 0 at threshold 0.5
        // Left child (node 1): leaf with value 0.2
        // Right child (node 2): leaf with value 0.9
        let children_left = vec![1, -1, -1];
        let children_right = vec![2, -1, -1];
        let features = vec![0, -2, -2];
        let thresholds = vec![0.5, -2.0, -2.0];
        let values = vec![0.0, 0.2, 0.9];

        // Sample with feature[0] = 0.3 <= 0.5 -> go left -> 0.2
        let sample_left = vec![0.3, 1.0];
        assert!((traverse_tree(&sample_left, &children_left, &children_right, &features, &thresholds, &values) - 0.2).abs() < 1e-10);

        // Sample with feature[0] = 0.7 > 0.5 -> go right -> 0.9
        let sample_right = vec![0.7, 1.0];
        assert!((traverse_tree(&sample_right, &children_left, &children_right, &features, &thresholds, &values) - 0.9).abs() < 1e-10);
    }
}
