"""From-scratch elastic-net penalized least squares via coordinate descent
-- the algorithm behind glmnet (Friedman, Hastie & Tibshirani, JSS 2010,
"Regularization Paths for Generalized Linear Models via Coordinate Descent").

Objective (scikit-learn parameterization, so we can validate against it exactly):

    (1/2n) ||y - Xb||^2  +  alpha*l1_ratio*||b||_1  +  (alpha*(1-l1_ratio)/2)*||b||^2

    l1_ratio = 1  ->  Lasso   (pure L1, sparse)
    l1_ratio = 0  ->  Ridge   (pure L2, shrinks, never zeros)
    0 < l1_ratio < 1 -> Elastic Net

The two penalties are the frequentist face of Bayesian shrinkage priors:
Ridge is the MAP estimate under an i.i.d. Gaussian prior on the coefficients,
Lasso the MAP under an i.i.d. Laplace (double-exponential) prior -- the same
priors that, treated fully Bayesianly, drive the Variable-Selection arc
(SSVS, horseshoe, BMA).  Kernel ridge (companion SVM notebook) is in turn the
posterior mean of a Gaussian process (BNP arc).

Coordinate descent cycles one coefficient at a time; with the other
coefficients fixed the 1-D penalized least-squares solution is the
soft-thresholding operator, which is what makes L1 produce exact zeros:

    b_j <- S(rho_j, alpha*l1_ratio) / (z_j + alpha*(1-l1_ratio)),
    rho_j = (1/n) x_j . (r + x_j b_j),   z_j = (1/n)||x_j||^2,
    S(z, g) = sign(z) * max(|z| - g, 0).
"""
import numpy as np


def soft_threshold(z, g):
    """Soft-thresholding operator S(z, g) -- the 1-D lasso solution."""
    return np.sign(z) * np.maximum(np.abs(z) - g, 0.0)


class ElasticNetCD:
    """Elastic-net regression by cyclic coordinate descent.

    Matches sklearn's ElasticNet objective (fit on centered X; intercept
    recovered as ybar - xbar.b), so coefficients agree to numerical
    tolerance.  Standardize the columns yourself before fitting for a fair,
    scale-free penalty (glmnet does this internally); the notebook does.
    """

    def __init__(self, alpha=1.0, l1_ratio=1.0, max_iter=1000, tol=1e-8):
        self.alpha = alpha
        self.l1_ratio = l1_ratio
        self.max_iter = max_iter
        self.tol = tol

    def fit(self, X, y):
        X = np.asarray(X, float); y = np.asarray(y, float)
        n, p = X.shape
        self.xm_ = X.mean(0); self.ym_ = y.mean()
        Xc = X - self.xm_; yc = y - self.ym_
        z = (Xc ** 2).sum(0) / n                       # per-column (1/n)||x_j||^2
        z[z == 0] = 1.0
        a1 = self.alpha * self.l1_ratio                # L1 strength
        a2 = self.alpha * (1 - self.l1_ratio)          # L2 strength
        b = np.zeros(p)
        r = yc - Xc @ b                                # residual
        for it in range(self.max_iter):
            dmax = 0.0; bmax = 0.0
            for j in range(p):
                bj = b[j]
                rho = Xc[:, j] @ (r + Xc[:, j] * bj) / n
                b[j] = soft_threshold(rho, a1) / (z[j] + a2)
                if b[j] != bj:
                    r += Xc[:, j] * (bj - b[j])        # keep residual in sync
                dmax = max(dmax, abs(b[j] - bj)); bmax = max(bmax, abs(b[j]))
            if dmax < self.tol * max(bmax, 1e-12):
                break
        self.coef_ = b
        self.intercept_ = self.ym_ - self.xm_ @ b
        self.n_iter_ = it + 1
        return self

    def predict(self, X):
        return np.asarray(X, float) @ self.coef_ + self.intercept_


def enet_path(X, y, l1_ratio=1.0, alphas=None, n_alphas=100, eps=1e-3,
              max_iter=1000, tol=1e-8):
    """Warm-started coefficient path over a grid of alpha (penalty) values.

    Returns (alphas, coefs) with coefs shape (p, n_alphas).  alpha_max is the
    smallest penalty that zeros every coefficient -- max_j |x_j.yc|/n / l1_ratio
    -- and the grid runs down from there on a log scale, each fit warm-started
    from the previous (the glmnet trick that makes whole paths cheap).
    """
    X = np.asarray(X, float); y = np.asarray(y, float)
    n, p = X.shape
    Xc = X - X.mean(0); yc = y - y.mean()
    l1 = max(l1_ratio, 1e-3)
    alpha_max = np.max(np.abs(Xc.T @ yc)) / (n * l1)
    if alphas is None:
        alphas = np.logspace(np.log10(alpha_max), np.log10(alpha_max * eps), n_alphas)
    coefs = np.zeros((p, len(alphas)))
    b = np.zeros(p)
    for k, a in enumerate(alphas):
        m = ElasticNetCD(alpha=a, l1_ratio=l1_ratio, max_iter=max_iter, tol=tol)
        m.coef_ = b.copy()                             # warm start (via a manual seed)
        # run coordinate descent seeded at b
        a1 = a * l1_ratio; a2 = a * (1 - l1_ratio)
        z = (Xc ** 2).sum(0) / n; z[z == 0] = 1.0
        r = yc - Xc @ b
        for _ in range(max_iter):
            dmax = 0.0; bmax = 0.0
            for j in range(p):
                bj = b[j]
                rho = Xc[:, j] @ (r + Xc[:, j] * bj) / n
                b[j] = soft_threshold(rho, a1) / (z[j] + a2)
                if b[j] != bj:
                    r += Xc[:, j] * (bj - b[j])
                dmax = max(dmax, abs(b[j] - bj)); bmax = max(bmax, abs(b[j]))
            if dmax < tol * max(bmax, 1e-12):
                break
        coefs[:, k] = b
    return alphas, coefs