onnx-export-rs 0.1.1

//! Version-pinned adapters for SmartCore models with private fitted state.
//!
//! These adapters use SmartCore's `serde` representation because SmartCore
//! 0.5.x does not expose the required state through public getters. They are
//! intentionally isolated behind `smartcore-compat`; serialization layout
//! drift is reported as [`crate::Error::InvalidModel`].

use ndarray::{Array1, Array2};
use serde::Serialize;
use serde_json::Value;

use crate::canonical::{
    flatten_tree, AffineModel, AggregationMode, CategoricalNaiveBayes, CentroidModel, DbscanModel,
    ForestStructure, GaussianNaiveBayes, GradientBoostedEnsemble, KnnClassifier, KnnRegressor,
    KnnWeight, LinearScoreClassifier, PostTransform, RecursiveNode, SvmClassifier, SvmKernel,
    SvmRegressor, TreeStructure, TreeTask,
};
use crate::{Error, Result};

fn incompatible(message: impl Into<String>) -> Error {
    Error::InvalidModel(format!(
        "incompatible SmartCore 0.5 serialization: {}",
        message.into()
    ))
}

fn serialized<T: Serialize>(model: &T) -> Result<Value> {
    serde_json::to_value(model).map_err(|error| incompatible(error.to_string()))
}

fn field<'a>(value: &'a Value, name: &str) -> Result<&'a Value> {
    value
        .get(name)
        .ok_or_else(|| incompatible(format!("missing `{name}`")))
}

fn some<'a>(value: &'a Value, name: &str) -> Result<&'a Value> {
    let value = field(value, name)?;
    if value.is_null() {
        Err(incompatible(format!("unfitted `{name}`")))
    } else {
        Ok(value)
    }
}

fn number(value: &Value, name: &str) -> Result<f64> {
    field(value, name)?
        .as_f64()
        .ok_or_else(|| incompatible(format!("invalid `{name}`")))
}

fn optional_number(value: &Value, name: &str) -> Result<Option<f64>> {
    let value = field(value, name)?;
    if value.is_null() {
        Ok(None)
    } else {
        value
            .as_f64()
            .map(Some)
            .ok_or_else(|| incompatible(format!("invalid `{name}`")))
    }
}

fn index(value: &Value, name: &str) -> Result<usize> {
    field(value, name)?
        .as_u64()
        .and_then(|value| usize::try_from(value).ok())
        .ok_or_else(|| incompatible(format!("invalid `{name}`")))
}

fn optional_index(value: &Value, name: &str) -> Result<Option<usize>> {
    let value = field(value, name)?;
    if value.is_null() {
        Ok(None)
    } else {
        value
            .as_u64()
            .and_then(|value| usize::try_from(value).ok())
            .map(Some)
            .ok_or_else(|| incompatible(format!("invalid `{name}`")))
    }
}

fn flat_nodes(nodes: &Value, class_count: Option<usize>) -> Result<TreeStructure> {
    let nodes = nodes
        .as_array()
        .ok_or_else(|| incompatible("invalid `nodes`"))?;
    let mut canonical = Vec::with_capacity(nodes.len());
    for (id, node) in nodes.iter().enumerate() {
        let left = optional_index(node, "true_child")?;
        let right = optional_index(node, "false_child")?;
        let leaf_values = if left.is_none() && right.is_none() {
            if let Some(class_count) = class_count {
                let output = index(node, "output")?;
                if output >= class_count {
                    return Err(incompatible("tree leaf class index is out of range"));
                }
                let mut scores = vec![0.0; class_count];
                scores[output] = 1.0;
                scores
            } else {
                vec![number(node, "output")?]
            }
        } else {
            Vec::new()
        };
        canonical.push(crate::canonical::TreeNode {
            id: i64::try_from(id).map_err(|_| incompatible("too many tree nodes"))?,
            feature_id: i64::try_from(index(node, "split_feature")?)
                .map_err(|_| incompatible("feature index is too large"))?,
            threshold: field(node, "split_value")?.as_f64().unwrap_or(0.0) as f32,
            true_child_id: i64::try_from(left.unwrap_or(0))
                .map_err(|_| incompatible("child index is too large"))?,
            false_child_id: i64::try_from(right.unwrap_or(0))
                .map_err(|_| incompatible("child index is too large"))?,
            branch_mode: crate::canonical::BranchMode::LessOrEqual,
            leaf_values: leaf_values.into_iter().map(|value| value as f32).collect(),
        });
    }
    Ok(TreeStructure { nodes: canonical })
}

/// Converts a fitted SmartCore decision-tree regressor.
pub fn decision_tree_regressor<T: Serialize>(model: &T) -> Result<ForestStructure> {
    let model = serialized(model)?;
    let tree = some(&model, "tree_regressor")?;
    Ok(ForestStructure {
        trees: vec![flat_nodes(field(tree, "nodes")?, None)?],
        aggregation: AggregationMode::Average,
        n_targets: 1,
    })
}

/// Converts a fitted SmartCore decision-tree classifier and its integer labels.
pub fn decision_tree_classifier<T: Serialize>(model: &T) -> Result<(ForestStructure, Vec<i64>)> {
    let model = serialized(model)?;
    let labels = integer_labels(field(&model, "classes")?)?;
    let tree = flat_nodes(field(&model, "nodes")?, Some(labels.len()))?;
    Ok((
        ForestStructure {
            trees: vec![tree],
            aggregation: AggregationMode::Average,
            n_targets: labels.len(),
        },
        labels,
    ))
}

/// Converts a fitted SmartCore random-forest regressor.
pub fn random_forest_regressor<T: Serialize>(model: &T) -> Result<ForestStructure> {
    let model = serialized(model)?;
    let forest = some(&model, "forest_regressor")?;
    let trees = some(forest, "trees")?
        .as_array()
        .ok_or_else(|| incompatible("invalid forest `trees`"))?
        .iter()
        .map(|tree| flat_nodes(field(tree, "nodes")?, None))
        .collect::<Result<Vec<_>>>()?;
    Ok(ForestStructure {
        trees,
        aggregation: AggregationMode::Average,
        n_targets: 1,
    })
}

/// Converts a fitted SmartCore Extra Trees regressor.
///
/// SmartCore stores Extra Trees in the same fitted forest representation as
/// random-forest regression, so both share the canonical conversion.
pub fn extra_trees_regressor<T: Serialize>(model: &T) -> Result<ForestStructure> {
    random_forest_regressor(model)
}

/// Converts a fitted SmartCore random-forest classifier and its integer labels.
pub fn random_forest_classifier<T: Serialize>(model: &T) -> Result<(ForestStructure, Vec<i64>)> {
    let model = serialized(model)?;
    let labels = integer_labels(some(&model, "classes")?)?;
    let trees = some(&model, "trees")?
        .as_array()
        .ok_or_else(|| incompatible("invalid forest `trees`"))?
        .iter()
        .map(|tree| flat_nodes(field(tree, "nodes")?, Some(labels.len())))
        .collect::<Result<Vec<_>>>()?;
    Ok((
        ForestStructure {
            trees,
            aggregation: AggregationMode::Average,
            n_targets: labels.len(),
        },
        labels,
    ))
}

fn integer_labels(value: &Value) -> Result<Vec<i64>> {
    value
        .as_array()
        .ok_or_else(|| incompatible("invalid class labels"))?
        .iter()
        .map(|label| {
            label
                .as_i64()
                .ok_or_else(|| incompatible("class label is not an integer"))
        })
        .collect()
}

fn recursive_xgboost(node: &Value) -> Result<RecursiveNode> {
    let left = field(node, "left")?;
    let right = field(node, "right")?;
    if left.is_null() && right.is_null() {
        return Ok(RecursiveNode::Leaf(vec![number(node, "value")?]));
    }
    if left.is_null() || right.is_null() {
        return Err(incompatible("XGBoost node has only one child"));
    }
    Ok(RecursiveNode::Branch {
        feature: index(node, "split_feature_idx")?,
        threshold: number(node, "threshold")?,
        left: Box::new(recursive_xgboost(left)?),
        right: Box::new(recursive_xgboost(right)?),
    })
}

/// Converts a fitted SmartCore `XGRegressor`.
pub fn xgboost_regressor<T: Serialize>(model: &T) -> Result<GradientBoostedEnsemble> {
    let model = serialized(model)?;
    let parameters = some(&model, "parameters")?;
    let trees = some(&model, "regressors")?
        .as_array()
        .ok_or_else(|| incompatible("invalid XGBoost `regressors`"))?
        .iter()
        .map(|tree| flatten_tree(&recursive_xgboost(tree)?))
        .collect::<Result<Vec<_>>>()?;
    Ok(GradientBoostedEnsemble {
        trees,
        base_values: vec![number(parameters, "base_score")?],
        learning_rate: number(parameters, "learning_rate")?,
        n_targets: 1,
        task: TreeTask::Regression,
        post_transform: PostTransform::None,
    })
}

/// Converts fitted SmartCore k-means centroids.
pub fn kmeans<T: Serialize>(model: &T) -> Result<CentroidModel> {
    let model = serialized(model)?;
    let rows = numeric_matrix(field(&model, "centroids")?)?;
    let columns = rows[0].len();
    let centroids =
        Array2::from_shape_vec((rows.len(), columns), rows.into_iter().flatten().collect())
            .map_err(|error| incompatible(error.to_string()))?;
    CentroidModel::new(centroids)
}

/// Converts a fitted SmartCore Gaussian Naive Bayes classifier.
pub fn gaussian_naive_bayes<T: Serialize>(model: &T) -> Result<GaussianNaiveBayes> {
    let model = serialized(model)?;
    let distribution = field(some(&model, "inner")?, "distribution")?;
    let mean_rows = numeric_matrix(field(distribution, "theta")?)?;
    let variance_rows = numeric_matrix(field(distribution, "var")?)?;
    let rows = mean_rows.len();
    let columns = mean_rows[0].len();
    let means = Array2::from_shape_vec((rows, columns), mean_rows.into_iter().flatten().collect())
        .map_err(|error| incompatible(error.to_string()))?;
    let variances = Array2::from_shape_vec(
        (rows, columns),
        variance_rows.into_iter().flatten().collect(),
    )
    .map_err(|error| incompatible(error.to_string()))?;
    GaussianNaiveBayes::new(
        means,
        variances,
        Array1::from(numeric_vector(field(distribution, "class_priors")?)?),
        integer_labels(field(distribution, "class_labels")?)?,
    )
}

fn naive_bayes_distribution(model: &Value) -> Result<&Value> {
    field(some(model, "inner")?, "distribution")
}

/// Converts a fitted SmartCore Multinomial Naive Bayes classifier.
pub fn multinomial_naive_bayes<T: Serialize>(model: &T) -> Result<LinearScoreClassifier> {
    let model = serialized(model)?;
    let distribution = naive_bayes_distribution(&model)?;
    let rows = numeric_matrix(field(distribution, "feature_log_prob")?)?;
    let classes = rows.len();
    let features = rows[0].len();
    let coefficients =
        Array2::from_shape_fn((features, classes), |(feature, class)| rows[class][feature]);
    let priors = numeric_vector(field(distribution, "class_priors")?)?;
    LinearScoreClassifier::new(
        coefficients,
        Array1::from_iter(priors.into_iter().map(f64::ln)),
        integer_labels(field(distribution, "class_labels")?)?,
        None,
    )
}

/// Converts a fitted SmartCore Bernoulli Naive Bayes classifier, including
/// its optional input binarization threshold.
pub fn bernoulli_naive_bayes<T: Serialize>(model: &T) -> Result<LinearScoreClassifier> {
    let model = serialized(model)?;
    let distribution = naive_bayes_distribution(&model)?;
    let log_probability = numeric_matrix(field(distribution, "feature_log_prob")?)?;
    let classes = log_probability.len();
    let features = log_probability[0].len();
    let mut bias = numeric_vector(field(distribution, "class_priors")?)?
        .into_iter()
        .map(f64::ln)
        .collect::<Vec<_>>();
    let coefficients = Array2::from_shape_fn((features, classes), |(feature, class)| {
        let log_p = log_probability[class][feature];
        let log_not_p = (-log_p.exp()).ln_1p();
        bias[class] += log_not_p;
        log_p - log_not_p
    });
    LinearScoreClassifier::new(
        coefficients,
        Array1::from(bias),
        integer_labels(field(distribution, "class_labels")?)?,
        optional_number(&model, "binarize")?,
    )
}

/// Converts a fitted SmartCore Categorical Naive Bayes classifier.
pub fn categorical_naive_bayes<T: Serialize>(model: &T) -> Result<CategoricalNaiveBayes> {
    let model = serialized(model)?;
    let distribution = naive_bayes_distribution(&model)?;
    let features = field(distribution, "coefficients")?
        .as_array()
        .ok_or_else(|| incompatible("invalid categorical coefficients"))?;
    let mut tables = Vec::with_capacity(features.len());
    for feature in features {
        let class_rows = numeric_matrix(feature)?;
        let classes = class_rows.len();
        let categories = class_rows[0].len();
        tables.push(Array2::from_shape_fn(
            (categories, classes),
            |(category, class)| class_rows[class][category],
        ));
    }
    CategoricalNaiveBayes::new(
        tables,
        Array1::from_iter(
            numeric_vector(field(distribution, "class_priors")?)?
                .into_iter()
                .map(f64::ln),
        ),
        integer_labels(field(distribution, "class_labels")?)?,
    )
}

/// Converts a fitted SmartCore `StandardScaler` into an affine transform.
pub fn standard_scaler<T: Serialize>(model: &T) -> Result<AffineModel> {
    let model = serialized(model)?;
    let means = numeric_vector(field(&model, "means")?)?;
    let standard_deviations = numeric_vector(field(&model, "stds")?)?;
    if means.len() != standard_deviations.len() || means.is_empty() {
        return Err(incompatible("invalid StandardScaler statistics"));
    }
    let parameters = field(&model, "parameters")?;
    let with_mean = field(parameters, "with_mean")?
        .as_bool()
        .ok_or_else(|| incompatible("invalid StandardScaler `with_mean`"))?;
    let with_std = field(parameters, "with_std")?
        .as_bool()
        .ok_or_else(|| incompatible("invalid StandardScaler `with_std`"))?;
    let scales = standard_deviations
        .into_iter()
        .map(|standard_deviation| {
            if with_std {
                1.0 / standard_deviation.max(f64::MIN_POSITIVE)
            } else {
                1.0
            }
        })
        .collect::<Vec<_>>();
    let matrix = Array2::from_shape_fn((means.len(), means.len()), |(row, column)| {
        if row == column {
            scales[row]
        } else {
            0.0
        }
    });
    let bias = Array1::from_iter(means.into_iter().zip(scales).map(|(mean, scale)| {
        if with_mean {
            -mean * scale
        } else {
            0.0
        }
    }));
    AffineModel::new(matrix, bias)
}

/// Converts a fitted SmartCore Euclidean DBSCAN model.
///
/// Callers must use the Euclidean distance because SmartCore's serialized
/// zero-sized metric is not identifiable.
pub fn dbscan<T: Serialize>(model: &T) -> Result<DbscanModel> {
    let model = serialized(model)?;
    let algorithm = field(&model, "knn_algorithm")?
        .as_object()
        .and_then(|object| object.values().next())
        .ok_or_else(|| incompatible("invalid DBSCAN neighbor index"))?;
    let rows = numeric_matrix(field(algorithm, "data")?)?;
    let columns = rows[0].len();
    let samples =
        Array2::from_shape_vec((rows.len(), columns), rows.into_iter().flatten().collect())
            .map_err(|error| incompatible(error.to_string()))?;
    DbscanModel::new(
        samples,
        integer_labels(field(&model, "cluster_labels")?)?,
        index(&model, "num_classes")?,
        number(&model, "eps")?,
    )
}

fn knn_state(model: &Value) -> Result<(Array2<f64>, usize, KnnWeight)> {
    let algorithm = some(model, "knn_algorithm")?
        .as_object()
        .and_then(|object| object.values().next())
        .ok_or_else(|| incompatible("invalid k-NN algorithm state"))?;
    let rows = numeric_matrix(field(algorithm, "data")?)?;
    let columns = rows[0].len();
    let samples =
        Array2::from_shape_vec((rows.len(), columns), rows.into_iter().flatten().collect())
            .map_err(|error| incompatible(error.to_string()))?;
    let weight = match some(model, "weight")?.as_str() {
        Some("Uniform") => KnnWeight::Uniform,
        Some("Distance") => KnnWeight::Distance,
        _ => return Err(incompatible("unknown k-NN weighting strategy")),
    };
    Ok((samples, index(model, "k")?, weight))
}

/// Converts a fitted SmartCore Euclidean k-NN regressor.
///
/// The serialized search index contains the original training samples. The
/// caller must only use this adapter with SmartCore's Euclidean distance,
/// because zero-sized custom distance types are not identifiable in serde.
pub fn knn_regressor<T: Serialize>(model: &T) -> Result<KnnRegressor> {
    let model = serialized(model)?;
    let (samples, k, weight) = knn_state(&model)?;
    KnnRegressor::new(
        samples,
        Array1::from(numeric_vector(some(&model, "y")?)?),
        k,
        weight,
    )
}

/// Converts a fitted SmartCore Euclidean k-NN classifier.
///
/// See [`knn_regressor`] for the distance-identification restriction.
pub fn knn_classifier<T: Serialize>(model: &T) -> Result<KnnClassifier> {
    let model = serialized(model)?;
    let (samples, k, weight) = knn_state(&model)?;
    KnnClassifier::new(
        samples,
        integer_labels(some(&model, "y")?)?,
        integer_labels(some(&model, "classes")?)?,
        k,
        weight,
    )
}

/// Converts a fitted SmartCore SVR. The kernel must match its fit parameters.
pub fn svm_regressor<T: Serialize>(model: &T, kernel: SvmKernel) -> Result<SvmRegressor> {
    let model = serialized(model)?;
    let rows = numeric_matrix(some(&model, "instances")?)?;
    let coefficients = numeric_vector(some(&model, "w")?)?;
    let feature_count = rows.first().map_or(0, Vec::len);
    let support_vectors = Array2::from_shape_vec(
        (rows.len(), feature_count),
        rows.into_iter().flatten().collect(),
    )
    .map_err(|error| incompatible(error.to_string()))?;
    SvmRegressor::new(
        support_vectors,
        Array1::from(coefficients),
        number(&model, "b")?,
        kernel,
        false,
    )
}

/// Converts a fitted binary SmartCore SVC.
///
/// The kernel must match the parameters used during fitting. SmartCore stores
/// signed dual coefficients but not support-vector class membership; for a
/// binary SVC the coefficient sign supplies that membership. Multiclass SVC is
/// intentionally rejected because its serialized wrapper does not retain a
/// stable ONNX pairwise layout.
pub fn svm_classifier<T: Serialize>(model: &T, kernel: SvmKernel) -> Result<SvmClassifier> {
    let model = serialized(model)?;
    let labels = integer_labels(some(&model, "classes")?)?;
    if labels.len() != 2 {
        return Err(incompatible("only binary SVC is supported"));
    }
    let rows = numeric_matrix(some(&model, "instances")?)?;
    let weights = numeric_vector(some(&model, "w")?)?;
    if rows.len() != weights.len() {
        return Err(incompatible("SVC support-vector/weight mismatch"));
    }
    let mut grouped = Vec::with_capacity(rows.len());
    let mut coefficients = Vec::with_capacity(rows.len());
    let mut counts = Vec::with_capacity(2);
    for class in [0, 1] {
        let before = grouped.len();
        for (row, &weight) in rows.iter().zip(&weights) {
            if (class == 0 && weight <= 0.0) || (class == 1 && weight > 0.0) {
                grouped.extend_from_slice(row);
                // ONNX's binary voting direction is opposite SmartCore's
                // positive-decision convention.
                coefficients.push(-weight);
            }
        }
        counts.push(grouped.len() / rows[0].len() - before / rows[0].len());
    }
    let support_vectors = Array2::from_shape_vec((rows.len(), rows[0].len()), grouped)
        .map_err(|error| incompatible(error.to_string()))?;
    let classifier = SvmClassifier {
        support_vectors,
        coefficients: Array1::from(coefficients),
        rho: Array1::from(vec![number(&model, "b")?]),
        vectors_per_class: counts,
        class_labels: labels,
        prob_a: Vec::new(),
        prob_b: Vec::new(),
        kernel,
    };
    classifier.validate()?;
    Ok(classifier)
}

fn numeric_vector(value: &Value) -> Result<Vec<f64>> {
    value
        .as_array()
        .ok_or_else(|| incompatible("invalid numeric vector"))?
        .iter()
        .map(|value| {
            value
                .as_f64()
                .ok_or_else(|| incompatible("invalid numeric value"))
        })
        .collect()
}

fn numeric_matrix(value: &Value) -> Result<Vec<Vec<f64>>> {
    let rows = value
        .as_array()
        .ok_or_else(|| incompatible("invalid numeric matrix"))?;
    let matrix = rows
        .iter()
        .map(numeric_vector)
        .collect::<Result<Vec<_>>>()?;
    let width = matrix.first().map_or(0, Vec::len);
    if width == 0 || matrix.iter().any(|row| row.len() != width) {
        return Err(incompatible("ragged or empty numeric matrix"));
    }
    Ok(matrix)
}