"""Standalone synthetic teaching example. Run locally with Python, not in the HTML."""
import json
import numpy as np
import pandas as pd
from scipy.special import expit

def clean(value):
    if isinstance(value, dict): return {str(k):clean(v) for k,v in value.items()}
    if isinstance(value,(list,tuple,np.ndarray)): return [clean(v) for v in value]
    if isinstance(value,(np.integer,)): return int(value)
    if isinstance(value,(float,np.floating)): return float(value) if np.isfinite(value) else None
    if isinstance(value,(np.bool_,)): return bool(value)
    return value

def print_result(result):
    print(json.dumps(clean(result),ensure_ascii=False,allow_nan=False))


def ols_fit(X, y):
    """X excludes intercept. Return intercept, slopes, rank of centered X."""
    X, y = np.asarray(X, float), np.asarray(y, float)
    if X.ndim != 2 or y.shape != (len(X),) or len(X) == 0:
        raise ValueError("Expected nonempty X(n,p) and y(n)")
    if not np.isfinite(X).all() or not np.isfinite(y).all():
        raise ValueError("Resolve nonfinite data first")
    mean = X.mean(axis=0)
    coef, _, rank, _ = np.linalg.lstsq(X-mean, y-y.mean(), rcond=None)
    return float(y.mean()-mean@coef), coef, int(rank)

def ridge_fit(X, y, lam):
    """SSE/(2n) + lam*||beta||^2/2; do not penalize intercept."""
    if not np.isfinite(lam) or lam < 0:
        raise ValueError("lam must be nonnegative and finite")
    if lam == 0:
        b, coef, _ = ols_fit(X, y)
        return b, coef
    X, y = np.asarray(X, float), np.asarray(y, float)
    ols_fit(X, y)  # Validate the input; not used as the regularized solution.
    mean = X.mean(0)
    xc, yc = X-mean, y-y.mean()
    coef = np.linalg.solve(xc.T@xc + len(X)*lam*np.eye(X.shape[1]), xc.T@yc)
    return float(y.mean()-mean@coef), coef

def soft_threshold(value, lam):
    return np.sign(value)*np.maximum(np.abs(value)-lam, 0)

def lasso_fit(X, y, lam, iterations=2000, tol=1e-9):
    """Centered cyclic coordinate descent; SSE/(2n)+lam*sum(abs(beta))."""
    if not np.isfinite(lam) or lam < 0:
        raise ValueError("Invalid penalty")
    X, y = np.asarray(X, float), np.asarray(y, float)
    ols_fit(X, y)
    mean, target_mean = X.mean(0), y.mean()
    xc, yc = X-mean, y-target_mean
    coef = np.zeros(X.shape[1]); residual = yc.copy()
    scale = np.mean(xc*xc, axis=0)
    converged = False
    for step in range(iterations):
        old = coef.copy()
        for j in range(len(coef)):
            residual += xc[:, j]*coef[j]
            c = xc[:, j]@residual/len(X)
            coef[j] = soft_threshold(c, lam)/scale[j] if scale[j] > 0 else 0
            residual -= xc[:, j]*coef[j]
        if np.max(np.abs(coef-old)) < tol:
            converged = True
            break
    return float(target_mean-mean@coef), coef, converged

if __name__ == "__main__":
    from sklearn.linear_model import Ridge, Lasso
    rng = np.random.default_rng(42)
    X = rng.normal(size=(60,3)); X[:,2] = .95*X[:,0]+.1*X[:,2]
    X = (X-X.mean(0))/X.std(0)
    y = 2 + X@np.array([1.5,0.,.5])+rng.normal(0,.2,len(X))
    lam=.05
    b,w = ridge_fit(X,y,lam)
    reference = Ridge(alpha=len(X)*lam).fit(X,y)
    np.testing.assert_allclose(w,reference.coef_,atol=1e-10)
    bl,wl,converged = lasso_fit(X,y,lam)
    lasso = Lasso(alpha=lam,tol=1e-10,max_iter=10000).fit(X,y)
    np.testing.assert_allclose(wl,lasso.coef_,atol=1e-7)
    print_result({"lambda":lam,"ridge_alpha":len(X)*lam,"ridge_coef":w,
        "lasso_coef":wl,"lasso_converged":converged,
        "ridge_reference_max_error":abs(w-reference.coef_).max(),
        "lasso_reference_max_error":abs(wl-lasso.coef_).max()})