//! Elastic Net / LASSO / Ridge regression via coordinate descent.
//!
//! Pure-Rust implementation of regularized linear regression with L1 (LASSO),
//! L2 (Ridge), and mixed (Elastic Net) penalties. Includes cross-validation
//! for hyperparameter selection.

use pyo3::prelude::*;

// ===========================================================================
// Result type
// ===========================================================================

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct ElasticNetResult {
    #[pyo3(get)]
    pub coefficients: Vec<f64>,
    #[pyo3(get)]
    pub intercept: f64,
    #[pyo3(get)]
    pub alpha: f64,
    #[pyo3(get)]
    pub l1_ratio: f64,
    #[pyo3(get)]
    pub nonzero_count: usize,
    #[pyo3(get)]
    pub r_squared: f64,
}

#[pymethods]
impl ElasticNetResult {
    fn __repr__(&self) -> String {
        format!(
            "ElasticNetResult(alpha={:.6}, l1_ratio={:.2}, nonzero={}, R²={:.4})",
            self.alpha, self.l1_ratio, self.nonzero_count, self.r_squared
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

#[inline]
fn soft_threshold(z: f64, gamma: f64) -> f64 {
    if z > gamma {
        z - gamma
    } else if z < -gamma {
        z + gamma
    } else {
        0.0
    }
}

fn mean(v: &[f64]) -> f64 {
    if v.is_empty() {
        return 0.0;
    }
    v.iter().sum::<f64>() / v.len() as f64
}

fn validate_xy(x: &[Vec<f64>], y: &[f64]) -> PyResult<(usize, usize)> {
    if x.is_empty() || y.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x and y must not be empty",
        ));
    }
    let n = x.len();
    if n != y.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x and y must have the same number of rows",
        ));
    }
    let p = x[0].len();
    if p == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must have at least one feature",
        ));
    }
    for row in x {
        if row.len() != p {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "all rows in x must have the same length",
            ));
        }
    }
    Ok((n, p))
}

/// Core coordinate descent for elastic net.
fn coordinate_descent(
    x: &[Vec<f64>],
    y: &[f64],
    alpha: f64,
    l1_ratio: f64,
    max_iter: usize,
    tol: f64,
) -> (Vec<f64>, f64) {
    let n = x.len();
    let p = x[0].len();
    let nf = n as f64;

    // Center y
    let y_mean = mean(y);
    let y_centered: Vec<f64> = y.iter().map(|&yi| yi - y_mean).collect();

    // Center x and compute column norms
    let mut x_means = vec![0.0; p];
    for j in 0..p {
        let mut s = 0.0;
        for i in 0..n {
            s += x[i][j];
        }
        x_means[j] = s / nf;
    }

    // Pre-compute x_centered columns and column squared sums
    let mut x_col: Vec<Vec<f64>> = vec![vec![0.0; n]; p];
    let mut col_sq: Vec<f64> = vec![0.0; p];
    for j in 0..p {
        for i in 0..n {
            let v = x[i][j] - x_means[j];
            x_col[j][i] = v;
            col_sq[j] += v * v;
        }
    }

    let mut beta = vec![0.0; p];
    let mut residual = y_centered.clone();

    let l1_penalty = alpha * l1_ratio;
    let l2_penalty = alpha * (1.0 - l1_ratio);

    for _iter in 0..max_iter {
        let mut max_change = 0.0_f64;

        for j in 0..p {
            if col_sq[j] < 1e-15 {
                continue;
            }

            let old_beta = beta[j];

            // Compute partial residual dot product
            let mut rho = 0.0;
            for i in 0..n {
                rho += x_col[j][i] * (residual[i] + old_beta * x_col[j][i]);
            }
            rho /= nf;

            // Coordinate update
            let denom = col_sq[j] / nf + l2_penalty;
            let new_beta = soft_threshold(rho, l1_penalty) / denom;

            if (new_beta - old_beta).abs() > 1e-15 {
                let delta = new_beta - old_beta;
                for i in 0..n {
                    residual[i] -= delta * x_col[j][i];
                }
                max_change = max_change.max((new_beta - old_beta).abs());
            }

            beta[j] = new_beta;
        }

        if max_change < tol {
            break;
        }
    }

    // Compute intercept
    let mut intercept = y_mean;
    for j in 0..p {
        intercept -= beta[j] * x_means[j];
    }

    (beta, intercept)
}

fn compute_r_squared(x: &[Vec<f64>], y: &[f64], beta: &[f64], intercept: f64) -> f64 {
    let n = x.len();
    let y_mean = mean(y);

    let mut ss_res = 0.0;
    let mut ss_tot = 0.0;
    for i in 0..n {
        let mut pred = intercept;
        for j in 0..beta.len() {
            pred += beta[j] * x[i][j];
        }
        ss_res += (y[i] - pred).powi(2);
        ss_tot += (y[i] - y_mean).powi(2);
    }

    if ss_tot < 1e-15 {
        return 0.0;
    }
    finite_or_zero(1.0 - ss_res / ss_tot)
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Fit an elastic net model using coordinate descent.
#[pyfunction]
#[pyo3(signature = (x, y, alpha=1.0, l1_ratio=0.5, max_iter=1000, tol=1e-6))]
pub fn elastic_net_fit(
    x: Vec<Vec<f64>>,
    y: Vec<f64>,
    alpha: f64,
    l1_ratio: f64,
    max_iter: usize,
    tol: f64,
) -> PyResult<ElasticNetResult> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let (_n, _p) = validate_xy(&x, &y)?;

    if alpha < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "alpha must be non-negative",
        ));
    }
    let l1_ratio = l1_ratio.clamp(0.0, 1.0);

    let (beta, intercept) = coordinate_descent(&x, &y, alpha, l1_ratio, max_iter, tol);
    let nonzero = beta.iter().filter(|&&b| b.abs() > 1e-10).count();
    let r_sq = compute_r_squared(&x, &y, &beta, intercept);

    Ok(ElasticNetResult {
        coefficients: beta,
        intercept: finite_or_zero(intercept),
        alpha,
        l1_ratio,
        nonzero_count: nonzero,
        r_squared: r_sq,
    })
}

/// Predict using a fitted elastic net model.
#[pyfunction]
pub fn elastic_net_predict(x: Vec<Vec<f64>>, result: &ElasticNetResult) -> PyResult<Vec<f64>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let mut predictions = Vec::with_capacity(x.len());
    for row in &x {
        if row.len() != result.coefficients.len() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "feature dimension mismatch",
            ));
        }
        let mut pred = result.intercept;
        for (j, &xj) in row.iter().enumerate() {
            pred += result.coefficients[j] * xj;
        }
        predictions.push(finite_or_zero(pred));
    }
    Ok(predictions)
}

/// Fit LASSO (L1 penalty only, l1_ratio=1.0).
#[pyfunction]
#[pyo3(signature = (x, y, alpha=1.0))]
pub fn lasso_fit(x: Vec<Vec<f64>>, y: Vec<f64>, alpha: f64) -> PyResult<ElasticNetResult> {
    elastic_net_fit(x, y, alpha, 1.0, 1000, 1e-6)
}

/// Fit Ridge regression (L2 penalty only, l1_ratio=0.0).
#[pyfunction]
#[pyo3(signature = (x, y, alpha=1.0))]
pub fn ridge_fit(x: Vec<Vec<f64>>, y: Vec<f64>, alpha: f64) -> PyResult<ElasticNetResult> {
    elastic_net_fit(x, y, alpha, 0.0, 1000, 1e-6)
}

/// Cross-validated elastic net: tries all combos of alphas x l1_ratios.
#[pyfunction]
#[pyo3(signature = (x, y, alphas=None, l1_ratios=None, n_folds=5))]
pub fn elastic_net_cv(
    x: Vec<Vec<f64>>,
    y: Vec<f64>,
    alphas: Option<Vec<f64>>,
    l1_ratios: Option<Vec<f64>>,
    n_folds: usize,
) -> PyResult<ElasticNetResult> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let (n, _p) = validate_xy(&x, &y)?;
    let n_folds = n_folds.max(2).min(n);

    let alphas = alphas.unwrap_or_else(|| vec![0.001, 0.01, 0.1, 1.0, 10.0]);
    let l1_ratios = l1_ratios.unwrap_or_else(|| vec![0.1, 0.5, 0.7, 0.9, 1.0]);

    let fold_size = n / n_folds;

    let mut best_score = f64::NEG_INFINITY;
    let mut best_alpha = alphas[0];
    let mut best_l1 = l1_ratios[0];

    for &alpha in &alphas {
        for &l1_ratio in &l1_ratios {
            let mut total_score = 0.0;
            let mut valid_folds = 0;

            for fold in 0..n_folds {
                let start = fold * fold_size;
                let end = if fold == n_folds - 1 { n } else { start + fold_size };

                // Split into train and validation
                let mut x_train = Vec::new();
                let mut y_train = Vec::new();
                let mut x_val = Vec::new();
                let mut y_val = Vec::new();

                for i in 0..n {
                    if i >= start && i < end {
                        x_val.push(x[i].clone());
                        y_val.push(y[i]);
                    } else {
                        x_train.push(x[i].clone());
                        y_train.push(y[i]);
                    }
                }

                if x_train.len() < 2 || x_val.is_empty() {
                    continue;
                }

                let (beta, intercept) =
                    coordinate_descent(&x_train, &y_train, alpha, l1_ratio, 500, 1e-5);
                let r2 = compute_r_squared(&x_val, &y_val, &beta, intercept);
                total_score += r2;
                valid_folds += 1;
            }

            if valid_folds > 0 {
                let avg_score = total_score / valid_folds as f64;
                if avg_score > best_score {
                    best_score = avg_score;
                    best_alpha = alpha;
                    best_l1 = l1_ratio;
                }
            }
        }
    }

    // Fit final model with best params on all data
    elastic_net_fit(x, y, best_alpha, best_l1, 1000, 1e-6)
}

// ===========================================================================
// Registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<ElasticNetResult>()?;
    m.add_function(wrap_pyfunction!(elastic_net_fit, m)?)?;
    m.add_function(wrap_pyfunction!(elastic_net_predict, m)?)?;
    m.add_function(wrap_pyfunction!(lasso_fit, m)?)?;
    m.add_function(wrap_pyfunction!(ridge_fit, m)?)?;
    m.add_function(wrap_pyfunction!(elastic_net_cv, m)?)?;
    Ok(())
}
