//! Inference for scikit-learn HistGradientBoosting heads exported as flat
//! tree arrays (src/ml/export.py). Mirrors sklearn's `_predictor.pyx`: NaN
//! follows `missing_go_to_left`, otherwise `x <= threshold` goes left; raw =
//! baseline + leaf values summed in tree order (sklearn's order, so the
//! regression output is bit-identical).

use anyhow::{Result, bail};
use serde::Deserialize;

/// One node, packed (24 bytes): `x` is the split threshold of an internal
/// node and the output of a leaf.
#[derive(Debug, Clone, Copy)]
struct Node {
    x: f64,
    left: u32,
    right: u32,
    feature: u16,
    leaf: bool,
    missing_left: bool,
}

/// The exported arrays, as in sklearn's predictor node table.
#[derive(Deserialize)]
struct RawTree {
    feature: Vec<usize>,
    threshold: Vec<Option<f64>>,
    missing_left: Vec<bool>,
    left: Vec<usize>,
    right: Vec<usize>,
    leaf: Vec<bool>,
    value: Vec<f64>,
}

#[derive(Debug)]
pub struct Tree {
    nodes: Vec<Node>,
}

impl<'de> Deserialize<'de> for Tree {
    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        use serde::de::Error;
        let r = RawTree::deserialize(d)?;
        let n = r.value.len();
        if n == 0
            || [
                r.feature.len(),
                r.threshold.len(),
                r.missing_left.len(),
                r.left.len(),
                r.right.len(),
                r.leaf.len(),
            ]
            .iter()
            .any(|&l| l != n)
        {
            return Err(D::Error::custom("tree arrays have different lengths"));
        }
        let nodes = (0..n)
            .map(|i| {
                let leaf = r.leaf[i];
                if !leaf && (r.left[i] >= n || r.right[i] >= n || r.feature[i] > u16::MAX as usize)
                {
                    return Err(D::Error::custom(format!(
                        "tree node {i} points outside the tree"
                    )));
                }
                Ok(Node {
                    // +inf thresholds ("missing vs everything") are exported as null
                    x: if leaf {
                        r.value[i]
                    } else {
                        r.threshold[i].unwrap_or(f64::INFINITY)
                    },
                    left: r.left[i] as u32,
                    right: r.right[i] as u32,
                    feature: r.feature[i] as u16,
                    leaf,
                    missing_left: r.missing_left[i],
                })
            })
            .collect::<Result<_, _>>()?;
        Ok(Tree { nodes })
    }
}

impl Tree {
    #[inline]
    fn predict(&self, x: &[f64]) -> f64 {
        let mut node = &self.nodes[0];
        while !node.leaf {
            let v = x[node.feature as usize];
            let go_left = if v.is_nan() {
                node.missing_left
            } else {
                v <= node.x
            };
            node = &self.nodes[if go_left { node.left } else { node.right } as usize];
        }
        node.x
    }

    fn validate(&self, n_features: usize) -> Result<()> {
        if let Some(i) = self
            .nodes
            .iter()
            .position(|n| !n.leaf && n.feature as usize >= n_features)
        {
            bail!("tree node {i} reads feature outside the feature vector");
        }
        Ok(())
    }
}

#[derive(Debug, Deserialize, PartialEq, Eq, Clone, Copy)]
#[serde(rename_all = "lowercase")]
pub enum Kind {
    Regressor,
    Classifier,
}

#[derive(Debug, Deserialize)]
pub struct Head {
    pub kind: Kind,
    baseline: f64,
    trees: Vec<Tree>,
}

impl Head {
    /// Regressor: the prediction. Classifier: P(class 1), as
    /// `predict_proba(X)[:, 1]` (scipy's `expit`).
    pub fn predict(&self, x: &[f64]) -> f64 {
        let mut raw = self.baseline;
        for t in &self.trees {
            raw += t.predict(x);
        }
        match self.kind {
            Kind::Regressor => raw,
            Kind::Classifier => 1.0 / (1.0 + (-raw).exp()),
        }
    }

    pub fn validate(&self, n_features: usize) -> Result<()> {
        self.trees.iter().try_for_each(|t| t.validate(n_features))
    }
}