//! Inference for scikit-learn HistGradientBoosting models exported as flat
//! tree arrays (see export/export_fixture.py). Mirrors sklearn's
//! `_predictor.pyx`: NaN follows `missing_go_to_left`, otherwise
//! `x <= threshold` goes left; raw = baseline + sum of leaf values, summed in
//! tree order (same order as sklearn, so results are bit-identical).

use serde::Deserialize;
use std::collections::BTreeMap;

fn thresholds<'de, D: serde::Deserializer<'de>>(d: D) -> Result<Vec<f64>, D::Error> {
    let raw: Vec<Option<f64>> = Deserialize::deserialize(d)?;
    Ok(raw
        .into_iter()
        .map(|t| t.unwrap_or(f64::INFINITY))
        .collect())
}

#[derive(Deserialize)]
pub struct Tree {
    feature: Vec<usize>,
    #[serde(deserialize_with = "thresholds")]
    threshold: Vec<f64>,
    missing_left: Vec<bool>,
    left: Vec<usize>,
    right: Vec<usize>,
    leaf: Vec<bool>,
    value: Vec<f64>,
}

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

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

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

impl Head {
    /// Regressor: the prediction. Classifier: P(class 1), like predict_proba[:, 1].
    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()),
        }
    }
}

#[derive(Deserialize)]
pub struct Bundle {
    pub station_code: String,
    pub model_version: String,
    pub regression_target: Option<String>,
    pub feature_names: Vec<String>,
    pub horizons: Vec<u32>,
    pub sigma: BTreeMap<String, f64>,
    pub thresholds: BTreeMap<String, f64>,
    pub heads: BTreeMap<String, Head>,
}