sklearn-rs 0.1.0

use ndarray::{Array1, Array2};
use crate::base::{Estimator, Predictor, validate_features, validate_target};
use crate::error::{Result, SklearnError};

/// 线性回归模型，类似于 scikit-learn 的 LinearRegression
#[derive(Debug, Clone)]
pub struct LinearRegression {
    pub fit_intercept: bool,
}

impl Default for LinearRegression {
    fn default() -> Self {
        Self {
            fit_intercept: true,
        }
    }
}

impl LinearRegression {
    pub fn new(fit_intercept: bool) -> Self {
        Self { fit_intercept }
    }
}

/// 训练好的线性模型
#[derive(Debug, Clone)]
pub struct LinearModel {
    pub coefficients: Array1<f64>,
    pub intercept: f64,
}

impl LinearModel {
    pub fn new(coefficients: Array1<f64>, intercept: f64) -> Self {
        Self {
            coefficients,
            intercept,
        }
    }
}

impl Predictor for LinearModel {
    type Input = Array2<f64>;
    type Output = Array1<f64>;
    
    fn predict(&self, x: &Self::Input) -> Result<Self::Output> {
        validate_features(x)?;
        
        if x.ncols() != self.coefficients.len() {
            return Err(SklearnError::ShapeMismatch {
                expected: format!("{} 个特征", self.coefficients.len()),
                actual: format!("{} 个特征", x.ncols()),
            });
        }
        
        let mut predictions = Array1::zeros(x.nrows());
        
        for (i, row) in x.rows().into_iter().enumerate() {
            let mut prediction = self.intercept;
            for (j, &coef) in self.coefficients.iter().enumerate() {
                prediction += coef * row[j];
            }
            predictions[i] = prediction;
        }
        
        Ok(predictions)
    }
}

impl Estimator for LinearRegression {
    type Input = Array2<f64>;
    type Target = Array1<f64>;
    type Model = LinearModel;
    
    fn fit(&self, x: &Self::Input, y: &Self::Target) -> Result<Self::Model> {
        validate_features(x)?;
        validate_target(y)?;
        
        if x.nrows() != y.len() {
            return Err(SklearnError::ShapeMismatch {
                expected: format!("{} 个样本", x.nrows()),
                actual: format!("{} 个目标值", y.len()),
            });
        }
        
        // 检查样本数量是否足够
        let n_samples = x.nrows();
        let n_features = x.ncols();
        
        if self.fit_intercept {
            if n_samples <= n_features + 1 {
                return self.fit_underdetermined_with_intercept(x, y);
            }
            self.fit_with_intercept(x, y)
        } else {
            if n_samples <= n_features {
                return self.fit_underdetermined_without_intercept(x, y);
            }
            self.fit_without_intercept(x, y)
        }
    }
}

impl LinearRegression {
    fn fit_with_intercept(&self, x: &Array2<f64>, y: &Array1<f64>) -> Result<LinearModel> {
        let n_samples = x.nrows();
        let n_features = x.ncols();
        
        // 构建设计矩阵 [1, X]
        let mut design_matrix = Array2::zeros((n_samples, n_features + 1));
        for i in 0..n_samples {
            design_matrix[(i, 0)] = 1.0; // 截距列
            for j in 0..n_features {
                design_matrix[(i, j + 1)] = x[(i, j)];
            }
        }
        
        // 计算 X^T X
        let xt_x = self.matrix_transpose_dot(&design_matrix, &design_matrix);
        
        // 计算 X^T y
        let xt_y = self.matrix_transpose_dot_vector(&design_matrix, y);
        
        // 解线性方程组 (X^T X) * coefficients = X^T y
        let all_coefficients = self.solve_linear_system(&xt_x, &xt_y)?;
        
        // 分离截距和系数
        let intercept = all_coefficients[0];
        let coefficients = all_coefficients.slice(ndarray::s![1..]).to_owned();
        
        Ok(LinearModel::new(coefficients, intercept))
    }
    
    fn fit_without_intercept(&self, x: &Array2<f64>, y: &Array1<f64>) -> Result<LinearModel> {
        // 计算 X^T X
        let xt_x = self.matrix_transpose_dot(x, x);
        
        // 计算 X^T y
        let xt_y = self.matrix_transpose_dot_vector(x, y);
        
        // 解线性方程组
        let coefficients = self.solve_linear_system(&xt_x, &xt_y)?;
        
        Ok(LinearModel::new(coefficients, 0.0))
    }
    
    // 处理欠定系统（样本数 <= 特征数 + 1）
    fn fit_underdetermined_with_intercept(&self, x: &Array2<f64>, y: &Array1<f64>) -> Result<LinearModel> {
        let n_samples = x.nrows();
        let n_features = x.ncols();
        
        println!("警告: 欠定系统，样本数({}) <= 特征数({}) + 1", n_samples, n_features);
        println!("使用最小范数解");
        
        // 使用伪逆或最小范数解
        // 这里简化处理：如果只有一个样本，直接返回平均值
        if n_samples == 1 {
            let intercept = y[0];
            let coefficients = Array1::zeros(n_features);
            return Ok(LinearModel::new(coefficients, intercept));
        }
        
        // 对于其他欠定情况，使用正规方程但添加小的正则化项避免奇异
        let mut design_matrix = Array2::zeros((n_samples, n_features + 1));
        for i in 0..n_samples {
            design_matrix[(i, 0)] = 1.0;
            for j in 0..n_features {
                design_matrix[(i, j + 1)] = x[(i, j)];
            }
        }
        
        let xt_x = self.matrix_transpose_dot(&design_matrix, &design_matrix);
        let xt_y = self.matrix_transpose_dot_vector(&design_matrix, y);
        
        // 添加小的正则化项避免奇异
        let mut regularized_xt_x = xt_x.clone();
        for i in 0..regularized_xt_x.nrows() {
            regularized_xt_x[(i, i)] += 1e-8;
        }
        
        let all_coefficients = self.solve_linear_system(&regularized_xt_x, &xt_y)?;
        
        let intercept = all_coefficients[0];
        let coefficients = all_coefficients.slice(ndarray::s![1..]).to_owned();
        
        Ok(LinearModel::new(coefficients, intercept))
    }
    
    // 处理欠定系统（无截距）
    fn fit_underdetermined_without_intercept(&self, x: &Array2<f64>, y: &Array1<f64>) -> Result<LinearModel> {
        let n_samples = x.nrows();
        let n_features = x.ncols();
        
        println!("警告: 欠定系统，样本数({}) <= 特征数({})", n_samples, n_features);
        println!("使用最小范数解");
        
        // 如果只有一个样本，返回该样本对应的解
        if n_samples == 1 {
            let mut coefficients = Array1::zeros(n_features);
            // 简单处理：使用第一个特征来拟合
            if n_features > 0 && x[(0, 0)].abs() > 1e-10 {
                coefficients[0] = y[0] / x[(0, 0)];
            }
            return Ok(LinearModel::new(coefficients, 0.0));
        }
        
        let xt_x = self.matrix_transpose_dot(x, x);
        let xt_y = self.matrix_transpose_dot_vector(x, y);
        
        // 添加小的正则化项避免奇异
        let mut regularized_xt_x = xt_x.clone();
        for i in 0..regularized_xt_x.nrows() {
            regularized_xt_x[(i, i)] += 1e-8;
        }
        
        let coefficients = self.solve_linear_system(&regularized_xt_x, &xt_y)?;
        
        Ok(LinearModel::new(coefficients, 0.0))
    }
    
    // 手动实现矩阵转置乘法：A^T * B
    fn matrix_transpose_dot(&self, a: &Array2<f64>, b: &Array2<f64>) -> Array2<f64> {
        let n = a.ncols(); // A^T 的行数
        let m = b.ncols(); // B 的列数
        let mut result = Array2::zeros((n, m));
        
        for i in 0..n {
            for j in 0..m {
                let mut sum = 0.0;
                for k in 0..a.nrows() {
                    sum += a[(k, i)] * b[(k, j)];
                }
                result[(i, j)] = sum;
            }
        }
        
        result
    }
    
    // 手动实现矩阵转置乘向量：A^T * v
    fn matrix_transpose_dot_vector(&self, a: &Array2<f64>, v: &Array1<f64>) -> Array1<f64> {
        let n = a.ncols();
        let mut result = Array1::zeros(n);
        
        for i in 0..n {
            let mut sum = 0.0;
            for k in 0..a.nrows() {
                sum += a[(k, i)] * v[k];
            }
            result[i] = sum;
        }
        
        result
    }
    
    fn solve_linear_system(&self, a: &Array2<f64>, b: &Array1<f64>) -> Result<Array1<f64>> {
        let n = a.nrows();
        
        // 使用高斯消元法求解线性方程组
        let mut augmented = Array2::zeros((n, n + 1));
        
        // 构建增广矩阵 [A | b]
        for i in 0..n {
            for j in 0..n {
                augmented[(i, j)] = a[(i, j)];
            }
            augmented[(i, n)] = b[i];
        }
        
        // 高斯消元法
        for i in 0..n {
            // 寻找主元
            let mut max_row = i;
            for k in i + 1..n {
                if augmented[(k, i)].abs() > augmented[(max_row, i)].abs() {
                    max_row = k;
                }
            }
            
            // 交换行
            if max_row != i {
                for j in 0..=n {
                    let temp = augmented[(i, j)];
                    augmented[(i, j)] = augmented[(max_row, j)];
                    augmented[(max_row, j)] = temp;
                }
            }
            
            // 主元为0，矩阵奇异
            if augmented[(i, i)].abs() < 1e-10 {
                return Err(SklearnError::FitFailed {
                    reason: "矩阵奇异，无法求解".to_string(),
                });
            }
            
            // 消元
            for k in i + 1..n {
                let factor = augmented[(k, i)] / augmented[(i, i)];
                for j in i..=n {
                    augmented[(k, j)] -= factor * augmented[(i, j)];
                }
            }
        }
        
        // 回代
        let mut x = Array1::zeros(n);
        for i in (0..n).rev() {
            x[i] = augmented[(i, n)];
            for j in i + 1..n {
                x[i] -= augmented[(i, j)] * x[j];
            }
            x[i] /= augmented[(i, i)];
        }
        
        Ok(x)
    }
}