datarust 0.6.5

//! Core traits shared by the numeric and categorical transformers.

use crate::error::{DatarustError, Result};
use crate::matrix::{Matrix, StrMatrix};

/// Shared base contract for every estimator in datarust.
///
/// `Estimator` is the common base for transformers and supervised predictors.
/// It deliberately does not prescribe input or output types; those belong to
/// the more specific [`Transformer`], [`Regressor`], and [`Classifier`] traits.
pub trait Estimator {}

/// A type-erased scalar parameter value, used by [`Params`] for hyperparameter
/// introspection.
///
/// This enables [`GridSearchCV`](crate::model_selection)-style hyperparameter
/// search without runtime reflection: estimators enumerate their tunable
/// parameters as `(name, value)` pairs that the search driver can read and
/// write back before each refit.
#[derive(Debug, Clone, PartialEq)]
pub enum ParamValue {
    /// A floating-point hyperparameter (e.g. `alpha`, `tol`).
    Float(f64),
    /// An integer hyperparameter (e.g. `n_clusters`, `max_iter`).
    Int(usize),
    /// A boolean hyperparameter (e.g. `fit_intercept`).
    Bool(bool),
}

/// Hyperparameter introspection for estimators that support tuning.
///
/// Estimators opt into this trait to expose their tunable hyperparameters as a
/// flat `(name, value)` map. This is the foundation for
/// [`GridSearchCV`](crate::model_selection)-style hyperparameter search: the
/// search driver reads the current parameters via [`get_params`](Params::get_params),
/// applies a candidate combination via [`set_params`](Params::set_params), then
/// refits and scores the model.
///
/// Not every estimator needs to implement `Params` — only those whose
/// hyperparameters should be searchable.
pub trait Params {
    /// Returns the current hyperparameters as `(name, value)` pairs.
    fn get_params(&self) -> Vec<(&'static str, ParamValue)>;

    /// Applies a named hyperparameter. Returns an error if `name` is unknown or
    /// `value` is the wrong type for that parameter.
    fn set_params(&mut self, name: &str, value: ParamValue) -> Result<()>;
}

/// Trait for numeric transformers operating on `Matrix -> Matrix`.
///
/// All scalers, decompositions, and other numeric transformers implement this
/// trait.  Call [`fit`](Transformer::fit) to learn parameters from training
/// data, then [`transform`](Transformer::transform) to apply the learned
/// transformation to new data.
///
/// ```rust
/// use datarust::scaler::StandardScaler;
/// use datarust::traits::Transformer;
/// use datarust::Matrix;
///
/// let x = Matrix::new(vec![
///     vec![1.0, 10.0],
///     vec![2.0, 20.0],
///     vec![3.0, 30.0],
/// ])?;
/// let mut s = StandardScaler::new();
/// let out = s.fit_transform(&x)?;
/// assert_eq!(out.ncols(), 2);
/// // Inverse supported on StandardScaler
/// let back = s.inverse_transform(&out)?;
/// for i in 0..3 {
///     for j in 0..2 {
///         assert!((back.get(i, j) - x.get(i, j)).abs() < 1e-9);
///     }
/// }
/// # Ok::<_, Box<dyn std::error::Error>>(())
/// ```
pub trait Transformer: Estimator {
    /// Name of the transformer, used for diagnostics.
    fn name(&self) -> &'static str;

    /// Fit the transformer on training data.
    fn fit(&mut self, x: &Matrix) -> Result<()>;

    /// Transform data using fitted parameters.
    fn transform(&self, x: &Matrix) -> Result<Matrix>;

    /// Convenience: fit then transform.
    fn fit_transform(&mut self, x: &Matrix) -> Result<Matrix> {
        self.fit(x)?;
        self.transform(x)
    }

    /// Fit the transformer with supervised target values when they are
    /// available. Unsupervised transformers ignore `y` by default. Supervised
    /// feature selectors override this method, allowing them to live safely in
    /// a [`SupervisedPipeline`](crate::pipeline::SupervisedPipeline).
    fn fit_with_target(&mut self, x: &Matrix, _y: &[f64]) -> Result<()> {
        self.fit(x)
    }

    /// Reverse the transformation, recovering an approximation of the
    /// original input.  Not all transformers support this; the default
    /// implementation returns an error.
    fn inverse_transform(&self, _x: &Matrix) -> Result<Matrix> {
        Err(DatarustError::InvalidInput(format!(
            "{} does not support inverse_transform",
            Transformer::name(self)
        )))
    }

    /// Whether the transformer has been fitted.
    fn is_fitted(&self) -> bool;
}

/// Trait for transformers that can report their output feature names.
///
/// `input_features` are the names of the columns fed to `fit`. When `None`,
/// implementations generate synthetic names (e.g. `x0`, `x1`).
pub trait FeatureNames {
    /// Returns the output feature names given the optional input feature names.
    fn feature_names_out(&self, input_features: Option<&[String]>) -> Vec<String>;
}

/// Helper: generate default input names `x0..x{n-1}`.
pub fn default_input_names(n: usize) -> Vec<String> {
    (0..n).map(|i| format!("x{}", i)).collect()
}

/// Trait for categorical (string) transformers operating on `StrMatrix -> Matrix`.
///
/// All categorical encoders that accept a [`StrMatrix`] and
/// produce a [`Matrix`] implement this trait.
pub trait CategoricalTransformer: Estimator {
    /// Human-readable name of the transformer.
    fn name(&self) -> &'static str;

    /// Fit the transformer on categorical training data.
    fn fit(&mut self, x: &StrMatrix) -> Result<()>;

    /// Transform categorical data using fitted parameters, returning a numeric
    /// [`Matrix`].
    fn transform(&self, x: &StrMatrix) -> Result<Matrix>;

    /// Convenience: fit then transform.
    fn fit_transform(&mut self, x: &StrMatrix) -> Result<Matrix> {
        self.fit(x)?;
        self.transform(x)
    }

    /// Reverse the transformation, recovering category strings from numeric
    /// codes.  Not all encoders support this; the default implementation
    /// returns an error.
    fn inverse_transform(&self, _y: &Matrix) -> Result<StrMatrix> {
        Err(DatarustError::InvalidInput(format!(
            "{} does not support inverse_transform",
            CategoricalTransformer::name(self)
        )))
    }

    /// Whether the transformer has been fitted.
    fn is_fitted(&self) -> bool;
}

/// Trait for supervised encoders that require target values during `fit`.
///
/// Input is a [`StrMatrix`] of categorical features and a slice of target
/// values `&[f64]`.  Output is a numeric [`Matrix`].
pub trait TargetTransformer: Estimator {
    /// Human-readable name of the transformer.
    fn name(&self) -> &'static str;

    /// Fit the transformer on categorical features and target values.
    fn fit(&mut self, x: &StrMatrix, y: &[f64]) -> Result<()>;

    /// Transform categorical data using fitted parameters.
    fn transform(&self, x: &StrMatrix) -> Result<Matrix>;

    /// Convenience: fit then transform.
    fn fit_transform(&mut self, x: &StrMatrix, y: &[f64]) -> Result<Matrix> {
        self.fit(x, y)?;
        self.transform(x)
    }

    /// Reverse the transformation.  Not all supervised encoders support this;
    /// the default implementation returns an error.
    fn inverse_transform(&self, _y: &Matrix) -> Result<StrMatrix> {
        Err(DatarustError::InvalidInput(format!(
            "{} does not support inverse_transform",
            TargetTransformer::name(self)
        )))
    }

    /// Whether the transformer has been fitted.
    fn is_fitted(&self) -> bool;
}

/// Shared fitting and prediction contract for supervised estimators operating
/// on numeric features and targets.
pub trait Predictor: Estimator {
    /// Fit the estimator on training features and target values.
    fn fit(&mut self, x: &Matrix, y: &[f64]) -> Result<()>;

    /// Predict one numeric value per input row. For classifiers this is the
    /// predicted class label; probability estimates are exposed separately by
    /// [`PredictProba`].
    fn predict(&self, x: &Matrix) -> Result<Vec<f64>>;

    /// Convenience: fit then predict on the same data.
    fn fit_predict(&mut self, x: &Matrix, y: &[f64]) -> Result<Vec<f64>> {
        self.fit(x, y)?;
        self.predict(x)
    }

    /// Whether the estimator has learned fitted state.
    fn is_fitted(&self) -> bool;
}

/// Trait for regression estimators operating on `Matrix` features and
/// `&[f64]` targets.
///
/// Regression models (e.g. [`LinearRegression`]) implement this trait.
/// Call [`fit`](Predictor::fit) to learn coefficients from training data and
/// targets, then [`predict`](Predictor::predict) to generate predictions for
/// new data.
///
/// [`LinearRegression`]: crate::linear_model::LinearRegression
///
/// ```rust
/// use datarust::linear_model::LinearRegression;
/// use datarust::traits::Predictor;
/// use datarust::Matrix;
///
/// let x = Matrix::new(vec![
///     vec![1.0],
///     vec![2.0],
///     vec![3.0],
///     vec![4.0],
/// ])?;
/// let y = vec![3.0, 5.0, 7.0, 9.0]; // y = 2x + 1
///
/// let mut model = LinearRegression::new();
/// model.fit(&x, &y)?;
/// let pred = model.predict(&x)?;
/// assert!((pred[0] - 3.0).abs() < 1e-9);
/// assert!((pred[3] - 9.0).abs() < 1e-9);
/// # Ok::<_, Box<dyn std::error::Error>>(())
/// ```
pub trait Regressor: Predictor {
    /// Name of the estimator, used for diagnostics.
    ///
    /// Kept on `Regressor` for backwards compatibility with code that uses
    /// regression estimators through this trait.
    fn name(&self) -> &'static str;

    /// R² (coefficient of determination) of the prediction.
    fn score(&self, x: &Matrix, y: &[f64]) -> Result<f64> {
        let prediction = Predictor::predict(self, x)?;
        crate::metrics::regression::r2_score(y, &prediction)
    }
}

/// Trait for classifiers that return one class label per input row.
pub trait Classifier: Predictor {
    /// Mean classification accuracy, mirroring sklearn's default classifier
    /// score.
    fn score(&self, x: &Matrix, y: &[f64]) -> Result<f64> {
        let prediction = self.predict(x)?;
        crate::metrics::classification::accuracy_score(y, &prediction)
    }
}

/// Trait for classifiers that expose a probability for every class.
///
/// The returned matrix is shaped `(n_samples, n_classes)` and follows the
/// estimator's class order. Binary [`LogisticRegression`](crate::LogisticRegression)
/// returns columns `[P(class=0), P(class=1)]`.
pub trait PredictProba: Classifier {
    /// Return per-class probability estimates.
    fn predict_proba(&self, x: &Matrix) -> Result<Matrix>;
}

/// Trait for 1-D label encoders that map `&[String]` to `Vec<usize>`.
///
/// Used to encode target labels for supervised learning.
pub trait LabelTransformer: Estimator {
    /// Human-readable name of the transformer.
    fn name(&self) -> &'static str;

    /// Fit on label data, learning the unique sorted classes.
    fn fit(&mut self, x: &[String]) -> Result<()>;

    /// Transform labels to integer indices.
    fn transform(&self, x: &[String]) -> Result<Vec<usize>>;

    /// Reverses the transformation: indices back to strings.
    fn inverse_transform(&self, x: &[usize]) -> Result<Vec<String>>;

    /// Convenience: fit then transform.
    fn fit_transform(&mut self, x: &[String]) -> Result<Vec<usize>> {
        self.fit(x)?;
        self.transform(x)
    }

    /// Whether the transformer has been fitted.
    fn is_fitted(&self) -> bool;
}

/// Trait for clustering estimators operating on `Matrix` features with no target.
///
/// Clustering is unsupervised: `fit` takes only `X`, and [`predict`](Clusterer::predict)
/// returns cluster indices (`Vec<usize>`), not the regression targets or class
/// labels used by [`Predictor`]. This mirrors `sklearn.cluster` estimators such
/// as `KMeans`, which expose `fit_predict` and `predict` returning integer labels.
///
/// Implementors provide [`fit`](Clusterer::fit) and [`predict`](Clusterer::predict);
/// the convenience methods have default implementations.
pub trait Clusterer: Estimator {
    /// Human-readable name of the estimator, used for diagnostics.
    fn name(&self) -> &'static str;

    /// Fit the estimator to the data, learning cluster structure.
    fn fit(&mut self, x: &Matrix) -> Result<()>;

    /// Assign each input row to its nearest fitted cluster, returning the
    /// cluster index per row.
    fn predict(&self, x: &Matrix) -> Result<Vec<usize>>;

    /// Convenience: fit, then return the cluster index assigned to each training
    /// row. Equivalent to `fit` followed by `predict` on the same data, but
    /// implementations may compute the labels more efficiently during fitting.
    fn fit_predict(&mut self, x: &Matrix) -> Result<Vec<usize>> {
        self.fit(x)?;
        self.predict(x)
    }

    /// Convenience: fit then transform. The default implementation returns a
    /// one-hot encoding of the cluster assignments (one column per cluster);
    /// implementations may override it to produce a different embedding (e.g.
    /// spectral clustering's affinity matrix).
    fn fit_transform(&mut self, x: &Matrix) -> Result<Matrix> {
        let labels = self.fit_predict(x)?;
        let k = self.n_clusters();
        let n = x.nrows();
        let mut data = vec![0.0; n * k];
        for (i, &label) in labels.iter().enumerate() {
            data[i * k + label] = 1.0;
        }
        Matrix::from_flat(n, k, data)
    }

    /// Number of clusters learned during `fit`. Available only after fitting.
    fn n_clusters(&self) -> usize;

    /// Whether the estimator has been fitted.
    fn is_fitted(&self) -> bool;
}

macro_rules! impl_estimator_from_transformer {
    ($($ty:path),+ $(,)?) => {
        $(
            impl Estimator for $ty {}
        )+
    };
}

impl_estimator_from_transformer!(
    crate::scaler::StandardScaler,
    crate::scaler::MinMaxScaler,
    crate::scaler::MaxAbsScaler,
    crate::scaler::RobustScaler,
    crate::scaler::Normalizer,
    crate::scaler::Binarizer,
    crate::scaler::KBinsDiscretizer,
    crate::scaler::QuantileTransformer,
    crate::scaler::PowerTransformer,
    crate::polynomial::PolynomialFeatures,
    crate::selection::VarianceThreshold,
    crate::selection::SelectKBest,
    crate::decomposition::PCA,
    crate::decomposition::TruncatedSVD,
    crate::imputer::SimpleImputer,
    crate::imputer::KnnImputer,
    crate::function_transformer::FunctionTransformer,
    crate::transformer_kind::TransformerKind,
    crate::pipeline::Pipeline,
);

macro_rules! impl_estimator_from_categorical_transformer {
    ($($ty:path),+ $(,)?) => {
        $(
            impl Estimator for $ty {}
        )+
    };
}

impl_estimator_from_categorical_transformer!(
    crate::encoder::OneHotEncoder,
    crate::encoder::OrdinalEncoder,
    crate::encoder::FrequencyEncoder,
    crate::categorical_kind::CategoricalTransformerKind,
);

macro_rules! impl_estimator_from_target_transformer {
    ($($ty:path),+ $(,)?) => {
        $(
            impl Estimator for $ty {}
        )+
    };
}

impl_estimator_from_target_transformer!(
    crate::encoder::TargetEncoder,
    crate::target_kind::TargetTransformerKind,
);

impl Estimator for crate::encoder::LabelEncoder {}

macro_rules! impl_estimator_from_clusterer {
    ($($ty:path),+ $(,)?) => {
        $(
            impl Estimator for $ty {}
        )+
    };
}

impl_estimator_from_clusterer!(crate::cluster::KMeans,);

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

    #[derive(Default)]
    struct NumericTransformer {
        fitted: bool,
    }

    impl Estimator for NumericTransformer {}

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

        fn fit(&mut self, _x: &Matrix) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn transform(&self, x: &Matrix) -> Result<Matrix> {
            Ok(x.clone())
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

    #[derive(Default)]
    struct CategoricalTransformerStub {
        fitted: bool,
    }

    impl Estimator for CategoricalTransformerStub {}

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

        fn fit(&mut self, _x: &StrMatrix) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn transform(&self, x: &StrMatrix) -> Result<Matrix> {
            Matrix::zeros(x.nrows(), x.ncols())
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

    #[derive(Default)]
    struct TargetTransformerStub {
        fitted: bool,
    }

    impl Estimator for TargetTransformerStub {}

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

        fn fit(&mut self, _x: &StrMatrix, _y: &[f64]) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn transform(&self, x: &StrMatrix) -> Result<Matrix> {
            Matrix::zeros(x.nrows(), x.ncols())
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

    #[derive(Default)]
    struct PredictorStub {
        fitted: bool,
    }

    impl Estimator for PredictorStub {}

    impl Predictor for PredictorStub {
        fn fit(&mut self, _x: &Matrix, _y: &[f64]) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn predict(&self, x: &Matrix) -> Result<Vec<f64>> {
            Ok((0..x.nrows()).map(|i| i as f64).collect())
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

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

    impl Classifier for PredictorStub {}

    #[derive(Default)]
    struct LabelTransformerStub {
        fitted: bool,
    }

    impl Estimator for LabelTransformerStub {}

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

        fn fit(&mut self, _x: &[String]) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn transform(&self, x: &[String]) -> Result<Vec<usize>> {
            Ok((0..x.len()).collect())
        }

        fn inverse_transform(&self, x: &[usize]) -> Result<Vec<String>> {
            Ok(x.iter().map(|i| i.to_string()).collect())
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

    #[derive(Default)]
    struct ClustererStub {
        fitted: bool,
    }

    impl Estimator for ClustererStub {}

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

        fn fit(&mut self, _x: &Matrix) -> Result<()> {
            self.fitted = true;
            Ok(())
        }

        fn predict(&self, x: &Matrix) -> Result<Vec<usize>> {
            Ok((0..x.nrows()).map(|i| i % 2).collect())
        }

        fn n_clusters(&self) -> usize {
            2
        }

        fn is_fitted(&self) -> bool {
            self.fitted
        }
    }

    fn numeric_input() -> Matrix {
        Matrix::new(vec![vec![1.0], vec![2.0]]).unwrap()
    }

    #[test]
    fn default_numeric_transformer_methods_delegate_and_report_unsupported_inverse() {
        let x = numeric_input();
        let mut transformer = NumericTransformer::default();

        let transformed = transformer.fit_transform(&x).unwrap();
        assert_eq!(transformed, x);
        assert!(transformer.is_fitted());
        transformer.fit_with_target(&x, &[0.0, 1.0]).unwrap();
        assert!(matches!(
            transformer.inverse_transform(&x),
            Err(DatarustError::InvalidInput(_))
        ));
        assert_eq!(
            default_input_names(3),
            vec!["x0".to_string(), "x1".to_string(), "x2".to_string()]
        );
    }

    #[test]
    fn default_categorical_and_target_methods_delegate_and_report_unsupported_inverse() {
        let x = StrMatrix::from_column(["a", "b"]).unwrap();
        let numeric = numeric_input();

        let mut categorical = CategoricalTransformerStub::default();
        assert_eq!(categorical.fit_transform(&x).unwrap().nrows(), 2);
        assert!(matches!(
            categorical.inverse_transform(&numeric),
            Err(DatarustError::InvalidInput(_))
        ));

        let mut target = TargetTransformerStub::default();
        assert_eq!(target.fit_transform(&x, &[0.0, 1.0]).unwrap().ncols(), 1);
        assert!(matches!(
            target.inverse_transform(&numeric),
            Err(DatarustError::InvalidInput(_))
        ));
    }

    #[test]
    fn default_predictor_regressor_classifier_and_label_helpers_delegate() {
        let x = numeric_input();
        let y = [0.0, 1.0];
        let mut predictor = PredictorStub::default();

        assert_eq!(predictor.fit_predict(&x, &y).unwrap(), y.to_vec());
        assert_eq!(Regressor::score(&predictor, &x, &y).unwrap(), 1.0);
        assert_eq!(Classifier::score(&predictor, &x, &y).unwrap(), 1.0);

        let labels = vec!["low".to_string(), "high".to_string()];
        let mut label_transformer = LabelTransformerStub::default();
        assert_eq!(
            label_transformer.fit_transform(&labels).unwrap(),
            vec![0, 1]
        );
    }

    #[test]
    fn default_clusterer_helpers_return_predictions_and_one_hot_encoding() {
        let x = numeric_input();
        let mut clusterer = ClustererStub::default();

        assert_eq!(clusterer.fit_predict(&x).unwrap(), vec![0, 1]);
        let transformed = clusterer.fit_transform(&x).unwrap();
        assert_eq!(transformed.rows_ref(), vec![vec![1.0, 0.0], vec![0.0, 1.0]]);
    }
}