onnx-export-rs 0.1.1

//! Version-pinned adapters for Linfa models whose fitted state is private.
//!
//! These functions target Linfa 0.8.1's serde layout. Layout drift is returned
//! as an invalid-model error rather than silently producing a wrong graph.

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

use crate::canonical::{
    AffineModel, GaussianNaiveBayes, GeneralizedLinearModel, LinearScoreClassifier, LinkFunction,
    SvmClassifier, SvmKernel, SvmRegressor,
};
use crate::{Error, Result};

fn invalid(message: impl Into<String>) -> Error {
    Error::InvalidModel(format!(
        "incompatible Linfa 0.8.1 layout: {}",
        message.into()
    ))
}

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

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

fn vector(value: &Value) -> Result<Vec<f64>> {
    value
        .as_array()
        .or_else(|| value.get("data").and_then(Value::as_array))
        .ok_or_else(|| invalid("expected numeric vector"))?
        .iter()
        .map(|item| item.as_f64().ok_or_else(|| invalid("non-numeric value")))
        .collect()
}

fn matrix(value: &Value) -> Result<(usize, usize, Vec<f64>)> {
    if let Some(dimensions) = value.get("dim").and_then(Value::as_array) {
        let rows = dimensions
            .first()
            .and_then(Value::as_u64)
            .ok_or_else(|| invalid("invalid matrix rows"))? as usize;
        let columns = dimensions
            .get(1)
            .and_then(Value::as_u64)
            .ok_or_else(|| invalid("invalid matrix columns"))? as usize;
        return Ok((rows, columns, vector(value)?));
    }
    let rows = value
        .as_array()
        .ok_or_else(|| invalid("expected numeric matrix"))?;
    let parsed = rows.iter().map(vector).collect::<Result<Vec<_>>>()?;
    let columns = parsed.first().map_or(0, Vec::len);
    Ok((
        parsed.len(),
        columns,
        parsed.into_iter().flatten().collect(),
    ))
}

fn class_entries(value: &Value) -> Result<Vec<(i64, &Value)>> {
    let mut entries = value
        .as_object()
        .ok_or_else(|| invalid("class_info is not an object"))?
        .iter()
        .map(|(label, info)| {
            label
                .parse::<i64>()
                .map(|label| (label, info))
                .map_err(|_| invalid("class label is not an integer"))
        })
        .collect::<Result<Vec<_>>>()?;
    entries.sort_by_key(|entry| entry.0);
    Ok(entries)
}

/// Converts Linfa Gaussian Naive Bayes with integer labels.
pub fn gaussian_naive_bayes<T: Serialize>(model: &T) -> Result<GaussianNaiveBayes> {
    let model = value(model)?;
    let entries = class_entries(field(&model, "class_info")?)?;
    let means = entries
        .iter()
        .map(|(_, info)| vector(field(info, "theta")?))
        .collect::<Result<Vec<_>>>()?;
    let variances = entries
        .iter()
        .map(|(_, info)| vector(field(info, "sigma")?))
        .collect::<Result<Vec<_>>>()?;
    let features = means.first().map_or(0, Vec::len);
    GaussianNaiveBayes::new(
        Array2::from_shape_vec(
            (entries.len(), features),
            means.into_iter().flatten().collect(),
        )
        .map_err(|error| invalid(error.to_string()))?,
        Array2::from_shape_vec(
            (entries.len(), features),
            variances.into_iter().flatten().collect(),
        )
        .map_err(|error| invalid(error.to_string()))?,
        Array1::from_iter(
            entries
                .iter()
                .map(|(_, info)| {
                    field(info, "prior").and_then(|value| {
                        value.as_f64().ok_or_else(|| invalid("invalid class prior"))
                    })
                })
                .collect::<Result<Vec<_>>>()?,
        ),
        entries.iter().map(|entry| entry.0).collect(),
    )
}

/// Converts Linfa Multinomial Naive Bayes with integer labels.
pub fn multinomial_naive_bayes<T: Serialize>(model: &T) -> Result<LinearScoreClassifier> {
    let model = value(model)?;
    let entries = class_entries(field(&model, "class_info")?)?;
    let rows = entries
        .iter()
        .map(|(_, info)| vector(field(info, "feature_log_prob")?))
        .collect::<Result<Vec<_>>>()?;
    let features = rows.first().map_or(0, Vec::len);
    let coefficients = Array2::from_shape_fn((features, entries.len()), |(feature, class)| {
        rows[class][feature]
    });
    LinearScoreClassifier::new(
        coefficients,
        Array1::from_iter(
            entries
                .iter()
                .map(|(_, info)| {
                    field(info, "prior").and_then(|value| {
                        value
                            .as_f64()
                            .map(f64::ln)
                            .ok_or_else(|| invalid("invalid class prior"))
                    })
                })
                .collect::<Result<Vec<_>>>()?,
        ),
        entries.iter().map(|entry| entry.0).collect(),
        None,
    )
}

/// Converts a fitted Linfa `TweedieRegressor` (generalized linear model) into
/// a canonical GLM. The private link function is recovered from the serialized
/// state and maps to the exporter's inverse-link activation.
pub fn tweedie_regressor<T: Serialize>(model: &T) -> Result<GeneralizedLinearModel> {
    let model = value(model)?;
    let coefficients = Array1::from(vector(field(&model, "coef")?)?);
    let intercept = field(&model, "intercept")?
        .as_f64()
        .ok_or_else(|| invalid("invalid intercept"))?;
    let link = match field(&model, "link")?.as_str() {
        Some("Identity") => LinkFunction::Identity,
        Some("Log") => LinkFunction::Log,
        Some("Logit") => LinkFunction::Logit,
        other => return Err(invalid(format!("unknown link function {other:?}"))),
    };
    Ok(GeneralizedLinearModel::new(coefficients, intercept, link))
}

/// Maps a serialized Linfa `KernelMethod` onto the ONNX-ML kernel enum.
///
/// Linfa's Gaussian kernel is `exp(-||x - y||^2 / eps)`, so its ONNX RBF
/// coefficient is `gamma = 1 / eps`. Its polynomial kernel is
/// `(<x, y> + c)^d` with an implicit unit gamma. The unit `Linear` variant is
/// rejected here because linear SVMs do not retain their support vectors; use
/// [`crate::adapters::linfa::linear_svm_score`] instead.
fn svm_kernel(kernel: &Value) -> Result<SvmKernel> {
    if kernel.as_str() == Some("Linear") {
        return Err(invalid(
            "linear-kernel SVM stores no support vectors; use linear_svm_score",
        ));
    }
    if let Some(eps) = kernel.get("Gaussian") {
        let eps = eps.as_f64().ok_or_else(|| invalid("invalid Gaussian width"))?;
        if eps <= 0.0 {
            return Err(invalid("non-positive Gaussian kernel width"));
        }
        return Ok(SvmKernel::Rbf { gamma: 1.0 / eps });
    }
    if let Some(parameters) = kernel.get("Polynomial").and_then(Value::as_array) {
        let coef0 = parameters
            .first()
            .and_then(Value::as_f64)
            .ok_or_else(|| invalid("invalid polynomial constant"))?;
        let degree = parameters
            .get(1)
            .and_then(Value::as_f64)
            .ok_or_else(|| invalid("invalid polynomial degree"))?;
        if !degree.is_finite() || degree < 0.0 || degree.fract() != 0.0 || degree > f64::from(u32::MAX)
        {
            return Err(invalid("non-integer polynomial degree"));
        }
        return Ok(SvmKernel::Polynomial {
            gamma: 1.0,
            coef0,
            degree: degree as u32,
        });
    }
    Err(invalid("unsupported SVM kernel"))
}

/// Returns the dual coefficients Linfa retains as support vectors.
///
/// Linfa keeps every training coefficient in `alpha` but stores only the rows
/// whose magnitude exceeds `100 * epsilon` as support vectors, so the same
/// threshold selects the coefficients that align with them.
fn support_coefficients(alpha: &Value) -> Result<Vec<f64>> {
    Ok(vector(alpha)?
        .into_iter()
        .filter(|coefficient| coefficient.abs() > 100.0 * f64::EPSILON)
        .collect())
}

/// Converts a fitted nonlinear-kernel Linfa `Svm` regressor.
///
/// Linfa predicts `sum_i alpha_i * K(sv_i, x) - rho`, which matches the ONNX-ML
/// `SVMRegressor` score `sum_i coeff_i * K(sv_i, x) + rho_onnx` with
/// `coeff_i = alpha_i` and `rho_onnx = -rho`. Only the Gaussian and polynomial
/// kernels are supported because linear SVMs do not serialize their support
/// vectors; use [`crate::adapters::linfa::linear_svm_score`] for those.
pub fn svm_regressor<T: Serialize>(model: &T) -> Result<SvmRegressor> {
    let model = value(model)?;
    let SvmSupport {
        kernel,
        rho,
        coefficients,
        rows,
        columns,
        data,
    } = svm_support(&model)?;
    let support_vectors =
        Array2::from_shape_vec((rows, columns), data).map_err(|error| invalid(error.to_string()))?;
    SvmRegressor::new(support_vectors, Array1::from(coefficients), -rho, kernel, false)
}

/// Parsed pieces shared by the nonlinear SVM regressor and classifier adapters.
struct SvmSupport {
    kernel: SvmKernel,
    /// Linfa's `rho`; the ONNX-ML score offset is its negation.
    rho: f64,
    /// Signed dual coefficient for each support vector.
    coefficients: Vec<f64>,
    rows: usize,
    columns: usize,
    /// Row-major support-vector matrix data.
    data: Vec<f64>,
}

/// Extracts the kernel, `rho`, signed coefficients, and support-vector matrix
/// common to the nonlinear SVM regressor and classifier adapters.
fn svm_support(model: &Value) -> Result<SvmSupport> {
    let kernel = svm_kernel(field(model, "kernel_method")?)?;
    let rho = field(model, "rho")?
        .as_f64()
        .ok_or_else(|| invalid("invalid rho"))?;
    let coefficients = support_coefficients(field(model, "alpha")?)?;
    let support = field(model, "sep_hyperplane")?
        .get("WeightedCombination")
        .ok_or_else(|| invalid("nonlinear SVM is missing its support vectors"))?;
    let (rows, columns, data) = matrix(support)?;
    if rows != coefficients.len() {
        return Err(invalid("support-vector and coefficient counts disagree"));
    }
    Ok(SvmSupport {
        kernel,
        rho,
        coefficients,
        rows,
        columns,
        data,
    })
}

/// Converts a fitted binary nonlinear-kernel Linfa `Svm` classifier.
///
/// Linfa stores signed dual coefficients — `+alpha` for the positive (`true`)
/// class and `-alpha` for the negative (`false`) class — and predicts the
/// positive class when `sum_i alpha_i * K(sv_i, x) - rho >= 0`. This maps onto
/// ONNX-ML's `SVMClassifier`, whose binary decision
/// `sum_i coeff_i * K(sv_i, x) + rho_onnx` selects the first class label when
/// positive, with `coeff_i = alpha_i` and `rho_onnx = -rho`. Support vectors
/// are grouped by class as ONNX-ML expects, positive coefficients first.
///
/// `positive_label` is emitted for Linfa's `true` class and `negative_label`
/// for its `false` class. Only the Gaussian and polynomial kernels are
/// supported, because linear-kernel models store no support vectors; use
/// [`crate::adapters::linfa::linear_svm_score`] for those.
pub fn svm_classifier<T: Serialize>(
    model: &T,
    positive_label: i64,
    negative_label: i64,
) -> Result<SvmClassifier> {
    if positive_label == negative_label {
        return Err(invalid("SVM classifier needs two distinct labels"));
    }
    let model = value(model)?;
    let SvmSupport {
        kernel,
        rho,
        coefficients,
        rows,
        columns,
        data,
    } = svm_support(&model)?;
    // ONNX-ML lays out support vectors class by class. Linfa's positive class
    // owns the positive coefficients, so it is emitted first.
    let mut ordered_vectors = Vec::with_capacity(data.len());
    let mut ordered_coefficients = Vec::with_capacity(coefficients.len());
    let mut counts = [0usize; 2];
    for (class, positive) in [true, false].into_iter().enumerate() {
        for (row, &coefficient) in coefficients.iter().enumerate() {
            if (coefficient > 0.0) == positive {
                ordered_vectors.extend_from_slice(&data[row * columns..(row + 1) * columns]);
                ordered_coefficients.push(coefficient);
                counts[class] += 1;
            }
        }
    }
    let support_vectors = Array2::from_shape_vec((rows, columns), ordered_vectors)
        .map_err(|error| invalid(error.to_string()))?;
    let classifier = SvmClassifier {
        support_vectors,
        coefficients: Array1::from(ordered_coefficients),
        rho: Array1::from(vec![-rho]),
        vectors_per_class: counts.to_vec(),
        class_labels: vec![positive_label, negative_label],
        prob_a: Vec::new(),
        prob_b: Vec::new(),
        kernel,
    };
    classifier.validate()?;
    Ok(classifier)
}

/// Converts a fitted Linfa PLS prediction model into one affine graph.
pub fn pls<T: Serialize>(model: &T) -> Result<AffineModel> {
    let model = value(model)?;
    let inner = model.get(0).unwrap_or(&model);
    let (inputs, outputs, coefficient_values) = matrix(field(inner, "coefficients")?)?;
    let coefficients = Array2::from_shape_vec((inputs, outputs), coefficient_values)
        .map_err(|error| invalid(error.to_string()))?;
    let x_mean = vector(field(inner, "x_mean")?)?;
    let x_std = vector(field(inner, "x_std")?)?;
    let y_mean = vector(field(inner, "y_mean")?)?;
    let matrix = Array2::from_shape_fn((inputs, outputs), |(input, output)| {
        coefficients[(input, output)] / x_std[input]
    });
    let bias = Array1::from_iter((0..outputs).map(|output| {
        y_mean[output]
            - (0..inputs)
                .map(|input| x_mean[input] * matrix[(input, output)])
                .sum::<f64>()
    }));
    AffineModel::new(matrix, bias)
}