//
//  ElasticNet.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

import Foundation

/// Elastic-net linear regression by cyclic coordinate descent, following
/// glmnet's conventions exactly so results are cross-checkable:
///
///   min (1/2n)·Σᵢ(yᵢ − b₀ − xᵢ'β)² + λ·[α‖β‖₁ + ((1−α)/2)‖β‖₂²]
///
/// Predictors are standardized internally (mean 0, variance 1 with the
/// 1/n denominator — glmnet's `standardize = TRUE`), the penalty applies
/// on the standardized scale, and coefficients are reported back on the
/// original scale with an unpenalized intercept.
public struct ZQElasticNetResult: Equatable, Sendable {
    public var coefficients: [(name: String, value: Double)]
    public var intercept: Double
    public var alpha: Double
    public var lambda: Double
    public var nonzeroCount: Int
    public var rSquared: Double
    public var iterations: Int

    public static func == (lhs: ZQElasticNetResult, rhs: ZQElasticNetResult) -> Bool {
        lhs.coefficients.elementsEqual(rhs.coefficients, by: { $0 == $1 })
            && lhs.intercept == rhs.intercept
            && lhs.alpha == rhs.alpha
            && lhs.lambda == rhs.lambda
            && lhs.nonzeroCount == rhs.nonzeroCount
            && lhs.rSquared == rhs.rSquared
            && lhs.iterations == rhs.iterations
    }
}

public enum ZQElasticNet {

    /// - Parameters:
    ///   - alpha: 1 = lasso, 0 = ridge, in between = elastic net.
    ///   - lambda: penalty strength on glmnet's scale.
    public static func fit(
        y: [Double],
        predictors: [(name: String, values: [Double])],
        alpha: Double,
        lambda: Double,
        maximumIterations: Int = 100_000,
        tolerance: Double = 1e-14
    ) throws -> ZQElasticNetResult {
        let n = y.count
        let p = predictors.count
        guard p > 0 else { throw ZQStatsError("elasticnet: regressors required") }
        guard n > 1 else { throw ZQStatsError("elasticnet: insufficient observations") }
        guard (0...1).contains(alpha) else {
            throw ZQStatsError("elasticnet: alpha() must be in [0, 1]")
        }
        guard lambda >= 0 else {
            throw ZQStatsError("elasticnet: lambda() must be nonnegative")
        }

        // Standardize (1/n variance, glmnet convention); center y.
        let meanY = y.reduce(0, +) / Double(n)
        var shifts = [Double](repeating: 0, count: p)
        var scales = [Double](repeating: 1, count: p)
        var standardized = [[Double]](repeating: [], count: p)
        for (j, column) in predictors.enumerated() {
            guard column.values.count == n else {
                throw ZQStatsError("regressor '\(column.name)' has wrong length")
            }
            let mean = column.values.reduce(0, +) / Double(n)
            let variance = column.values.reduce(0) {
                $0 + ($1 - mean) * ($1 - mean)
            } / Double(n)
            guard variance > 0 else {
                throw ZQStatsError("elasticnet: '\(column.name)' is constant")
            }
            shifts[j] = mean
            scales[j] = variance.squareRoot()
            standardized[j] = column.values.map { ($0 - mean) / scales[j] }
        }
        let centeredY = y.map { $0 - meanY }

        // Cyclic coordinate descent with residual updates. On the
        // standardized scale each column has (1/n)Σx² = 1, so
        //   β_j ← S(ρ_j, λα) / (1 + λ₂),  ρ_j = (1/n)x_j'r + β_j.
        //
        // glmnet quirk, matched deliberately: for gaussian models glmnet
        // standardizes y internally, which leaves the L1 penalty invariant
        // but scales the effective ridge penalty by 1/sd(y) — so
        // λ₂ = λ(1−α)/sd(y) (1/n variance). Ridge solutions therefore
        // depend on the scale of y, exactly as glmnet's do.
        let varY = centeredY.reduce(0) { $0 + $1 * $1 } / Double(n)
        guard varY > 0 else {
            throw ZQStatsError("elasticnet: the response is constant")
        }
        var beta = [Double](repeating: 0, count: p)
        var residuals = centeredY
        let l1 = lambda * alpha
        let l2 = lambda * (1 - alpha) / varY.squareRoot()
        var iterations = 0

        for iteration in 1...maximumIterations {
            iterations = iteration
            var maxChange = 0.0
            for j in 0..<p {
                var rho = beta[j]
                let column = standardized[j]
                var dot = 0.0
                for i in 0..<n { dot += column[i] * residuals[i] }
                rho += dot / Double(n)

                let updated = softThreshold(rho, l1) / (1 + l2)
                let change = updated - beta[j]
                if change != 0 {
                    for i in 0..<n { residuals[i] -= change * column[i] }
                    beta[j] = updated
                    maxChange = max(maxChange, abs(change))
                }
            }
            if maxChange < tolerance * max(1, beta.map(abs).max() ?? 1) {
                break
            }
            if iteration == maximumIterations {
                throw ZQStatsError("elasticnet: did not converge — increase lambda or iterations")
            }
        }

        // Back to the original scale.
        var coefficients: [(name: String, value: Double)] = []
        var intercept = meanY
        var nonzero = 0
        for j in 0..<p {
            let value = beta[j] / scales[j]
            coefficients.append((predictors[j].name, value))
            intercept -= shifts[j] * value
            if value != 0 { nonzero += 1 }
        }

        // Fit on the original data.
        var rss = 0.0
        var tss = 0.0
        for i in 0..<n {
            var fitted = intercept
            for j in 0..<p { fitted += coefficients[j].value * predictors[j].values[i] }
            rss += (y[i] - fitted) * (y[i] - fitted)
            tss += (y[i] - meanY) * (y[i] - meanY)
        }

        return ZQElasticNetResult(
            coefficients: coefficients,
            intercept: intercept,
            alpha: alpha,
            lambda: lambda,
            nonzeroCount: nonzero,
            rSquared: tss > 0 ? 1 - rss / tss : .nan,
            iterations: iterations
        )
    }

    /// The smallest λ at which every coefficient is zero (for α > 0):
    /// λ_max = max_j |(1/n)·x_j'y| / α on the standardized scale.
    public static func lambdaMax(
        y: [Double], predictors: [(name: String, values: [Double])], alpha: Double
    ) -> Double {
        let n = y.count
        let meanY = y.reduce(0, +) / Double(n)
        var best = 0.0
        for column in predictors {
            let mean = column.values.reduce(0, +) / Double(n)
            let variance = column.values.reduce(0) {
                $0 + ($1 - mean) * ($1 - mean)
            } / Double(n)
            guard variance > 0 else { continue }
            let scale = variance.squareRoot()
            var dot = 0.0
            for i in 0..<n {
                dot += (column.values[i] - mean) / scale * (y[i] - meanY)
            }
            best = max(best, abs(dot) / Double(n))
        }
        return alpha > 0 ? best / alpha : best * 1000
    }

    private static func softThreshold(_ value: Double, _ threshold: Double) -> Double {
        if value > threshold { return value - threshold }
        if value < -threshold { return value + threshold }
        return 0
    }
}