millwright 0.1.1

//! The smartcore backend — Millwright's first engine.
//!
//! Wraps a couple of smartcore models behind the framework's traits. The
//! `Frame -> DenseMatrix` conversion happens here, at the edge, via
//! [`as_dense`]; nothing above this module ever names a smartcore type.
//!
//! Phase 0 adapts two models to prove the contract across a real backend:
//! - [`RandomForest`] — the classifier from the design brief's API example.
//! - [`LinearRegression`] — the regression path.

use std::sync::Arc;

use smartcore::ensemble::random_forest_classifier::{
    RandomForestClassifier, RandomForestClassifierParameters,
};
use smartcore::linalg::basic::matrix::DenseMatrix;
use smartcore::linear::linear_regression::{
    LinearRegression as ScLinearRegression, LinearRegressionParameters,
};

use crate::error::{Error, Result};
use crate::frame::{Dataset, Frame};
use crate::traits::{Estimator, ParamValue, Predictor};

// The concrete smartcore model types, spelled once.
type ScForest = RandomForestClassifier<f64, i64, DenseMatrix<f64>, Vec<i64>>;
type ScLinReg = ScLinearRegression<f64, f64, DenseMatrix<f64>, Vec<f64>>;

/// Convert a [`Frame`] to smartcore's native `DenseMatrix<f64>`.
///
/// This is the one place the boundary type crosses into an engine's world.
pub fn as_dense(frame: &Frame) -> Result<DenseMatrix<f64>> {
    DenseMatrix::from_2d_vec(&frame.as_rows())
        .map_err(|e| Error::Backend(format!("DenseMatrix conversion failed: {e}")))
}

/// A random-forest classifier backed by smartcore.
///
/// Class labels are the (integral) values of the [`Dataset`] target, cast to
/// integers for smartcore and back to `f64` on predict.
#[derive(Clone)]
pub struct RandomForest {
    n_trees: u16,
    max_depth: Option<u16>,
    // Arc so the wrapper is cheaply Clone (the fitted model is immutable and
    // smartcore's own type is not Clone); a re-fit replaces the Arc.
    model: Option<Arc<ScForest>>,
}

impl RandomForest {
    /// A forest with smartcore's defaults (100 trees, unbounded depth).
    pub fn new() -> Self {
        RandomForest {
            n_trees: 100,
            max_depth: None,
            model: None,
        }
    }

    /// Set the number of trees.
    pub fn n_trees(mut self, n: u16) -> Self {
        self.n_trees = n;
        self
    }

    /// Set the maximum tree depth.
    pub fn max_depth(mut self, d: u16) -> Self {
        self.max_depth = Some(d);
        self
    }
}

impl Default for RandomForest {
    fn default() -> Self {
        RandomForest::new()
    }
}

impl Estimator for RandomForest {
    fn name(&self) -> &'static str {
        "RandomForest"
    }

    fn fit(&mut self, dataset: &Dataset) -> Result<()> {
        let x = as_dense(dataset.features())?;
        let y: Vec<i64> = dataset.target().iter().map(|v| v.round() as i64).collect();

        let mut params = RandomForestClassifierParameters::default().with_n_trees(self.n_trees);
        if let Some(d) = self.max_depth {
            params = params.with_max_depth(d);
        }

        let model = RandomForestClassifier::fit(&x, &y, params)
            .map_err(|e| Error::Backend(format!("RandomForest fit failed: {e}")))?;
        self.model = Some(Arc::new(model));
        Ok(())
    }

    fn set_param(&mut self, name: &str, value: ParamValue) -> Result<()> {
        match name {
            "n_trees" => self.n_trees = value.as_i64()? as u16,
            "max_depth" => self.max_depth = Some(value.as_i64()? as u16),
            other => {
                return Err(Error::Param(format!(
                    "RandomForest has no parameter '{other}'"
                )))
            }
        }
        Ok(())
    }

    #[cfg(feature = "onnx")]
    fn to_onnx_proto(&self) -> Result<onnx_export_rs::proto::ModelProto> {
        crate::onnx::ExportOnnx::to_onnx(self)
    }
}

impl Predictor for RandomForest {
    fn predict(&self, frame: &Frame) -> Result<Vec<f64>> {
        let model = self
            .model
            .as_ref()
            .ok_or_else(|| Error::NotFitted("RandomForest::predict".into()))?;
        let x = as_dense(frame)?;
        let y = model
            .predict(&x)
            .map_err(|e| Error::Backend(format!("RandomForest predict failed: {e}")))?;
        Ok(y.into_iter().map(|c| c as f64).collect())
    }
}

/// Ordinary least squares, backed by smartcore.
#[derive(Clone)]
pub struct LinearRegression {
    model: Option<Arc<ScLinReg>>,
}

impl LinearRegression {
    pub fn new() -> Self {
        LinearRegression { model: None }
    }
}

impl Default for LinearRegression {
    fn default() -> Self {
        LinearRegression::new()
    }
}

impl Estimator for LinearRegression {
    fn name(&self) -> &'static str {
        "LinearRegression"
    }

    fn fit(&mut self, dataset: &Dataset) -> Result<()> {
        let x = as_dense(dataset.features())?;
        let y: Vec<f64> = dataset.target().to_vec();
        let model = ScLinearRegression::fit(&x, &y, LinearRegressionParameters::default())
            .map_err(|e| Error::Backend(format!("LinearRegression fit failed: {e}")))?;
        self.model = Some(Arc::new(model));
        Ok(())
    }

    #[cfg(feature = "onnx")]
    fn to_onnx_proto(&self) -> Result<onnx_export_rs::proto::ModelProto> {
        crate::onnx::ExportOnnx::to_onnx(self)
    }
}

impl Predictor for LinearRegression {
    fn predict(&self, frame: &Frame) -> Result<Vec<f64>> {
        let model = self
            .model
            .as_ref()
            .ok_or_else(|| Error::NotFitted("LinearRegression::predict".into()))?;
        let x = as_dense(frame)?;
        model
            .predict(&x)
            .map_err(|e| Error::Backend(format!("LinearRegression predict failed: {e}")))
    }
}

/// Export a fitted [`RandomForest`] to an ONNX tree-ensemble classifier.
///
/// Uses onnx-export-rs's serde-based compat adapter, so it is independent of the
/// smartcore version onnx-export itself was built against.
#[cfg(feature = "onnx")]
impl crate::onnx::ExportOnnx for RandomForest {
    fn to_onnx(&self) -> Result<onnx_export_rs::proto::ModelProto> {
        use onnx_export_rs::adapters::smartcore_compat::random_forest_classifier;
        use onnx_export_rs::canonical::TreeTask;
        use onnx_export_rs::exporters::export_tree_ensemble;

        let model = self
            .model
            .as_ref()
            .ok_or_else(|| Error::NotFitted("RandomForest::to_onnx".into()))?;
        let (forest, _labels) = random_forest_classifier(&**model)
            .map_err(|e| Error::Backend(format!("RandomForest ONNX adapter failed: {e}")))?;
        export_tree_ensemble(&forest, TreeTask::Classification)
            .map_err(|e| Error::Backend(format!("RandomForest ONNX export failed: {e}")))
    }
}

/// Export a fitted [`LinearRegression`] to an ONNX `Gemm` graph.
///
/// The canonical weights are read from smartcore's own coefficient/intercept
/// accessors, so this is independent of onnx-export's smartcore version. Unlike
/// the tree ensemble, the resulting graph runs in tract.
#[cfg(feature = "onnx")]
impl crate::onnx::ExportOnnx for LinearRegression {
    fn to_onnx(&self) -> Result<onnx_export_rs::proto::ModelProto> {
        use onnx_export_rs::canonical::LinearModelWeights;
        use onnx_export_rs::exporters::export_linear;
        use smartcore::linalg::basic::arrays::Array;

        let model = self
            .model
            .as_ref()
            .ok_or_else(|| Error::NotFitted("LinearRegression::to_onnx".into()))?;
        let coef = model.coefficients();
        let (nr, nc) = coef.shape();
        let mut coefficients = Vec::with_capacity(nr * nc);
        for r in 0..nr {
            for c in 0..nc {
                coefficients.push(*coef.get((r, c)));
            }
        }
        let weights =
            LinearModelWeights::new(ndarray::Array1::from(coefficients), *model.intercept());
        Ok(export_linear(&weights))
    }
}

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

    #[test]
    fn random_forest_separates_two_clusters() {
        // Two well-separated classes on a single feature.
        let x = Frame::from_rows(
            vec![
                vec![0.0, 0.0],
                vec![0.5, 0.2],
                vec![0.1, 0.4],
                vec![9.0, 9.0],
                vec![9.5, 8.8],
                vec![8.9, 9.3],
            ],
            vec!["a".into(), "b".into()],
        )
        .unwrap();
        let ds = Dataset::new(x.clone(), vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0]).unwrap();

        let mut rf = RandomForest::new().n_trees(50);
        rf.fit(&ds).unwrap();

        let test = Frame::from_rows(
            vec![vec![0.2, 0.1], vec![9.2, 9.1]],
            vec!["a".into(), "b".into()],
        )
        .unwrap();
        assert_eq!(rf.predict(&test).unwrap(), vec![0.0, 1.0]);
    }

    #[test]
    fn linear_regression_recovers_a_line() {
        // y = 2x + 1
        let x = Frame::from_rows(
            vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0]],
            vec!["x".into()],
        )
        .unwrap();
        let ds = Dataset::new(x, vec![1.0, 3.0, 5.0, 7.0]).unwrap();
        let mut lr = LinearRegression::new();
        lr.fit(&ds).unwrap();

        let test = Frame::from_rows(vec![vec![4.0]], vec!["x".into()]).unwrap();
        let pred = lr.predict(&test).unwrap()[0];
        assert!((pred - 9.0).abs() < 1e-6, "expected ~9.0, got {pred}");
    }

    #[test]
    fn predict_before_fit_errors() {
        let f = Frame::from_rows(vec![vec![1.0]], vec!["x".into()]).unwrap();
        assert!(RandomForest::new().predict(&f).is_err());
        assert!(LinearRegression::new().predict(&f).is_err());
    }
}