"""Builds lasso_coordinate_descent.ipynb: a lecture demo (not a graded HW) companion
to REGRESSION_ADVANCED_claude_0910.tex, S4 (Lasso, More Carefully).

Not a fill-in-the-blank exercise -- fully worked, meant to be run once and read/shown
in class, same spirit as the lecture note's own worked examples and plots.
"""
import nbformat as nbf

nb = nbf.v4.new_notebook()
cells = []

cells.append(nbf.v4.new_markdown_cell(
"""# Lasso Coordinate Descent — Lecture Demo

Companion to `REGRESSION_ADVANCED_claude_0910.pdf`, Section 4 ("Lasso, More Carefully").
Not a graded exercise — a from-scratch implementation you can run and read alongside the
note, checked against `sklearn`, on the real Housing dataset used throughout HW1.

1. Reproduce the note's hand-computed 5-point numeric example with real code.
2. Full coordinate-descent Lasso from scratch on all 13 Housing features, checked against `sklearn.linear_model.Lasso`.
3. The regularization path: watch coefficients get driven to exactly zero, one at a time, as the penalty grows."""
))

cells.append(nbf.v4.new_code_cell(
"""import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Lasso

DATA = "../../data"
np.set_printoptions(suppress=True)"""
))

cells.append(nbf.v4.new_markdown_cell(
"""## Part 1 — the note's numeric example, reproduced from real data

The note's "unstable 5-point RM→MEDV fit" is the first 5 rows of the actual Housing training
set (feature `RM`, target `MEDV`). Same numbers as the note: $a^1=9.174$, $\\|\\mathbf{x}^1\\|^2=0.4808$,
$w^1_{OLS}=19.08$, $w^1_{ridge}(\\lambda=1)=6.20$, and Lasso hits exactly $0$ once $\\lambda_1$ crosses $20.80$."""
))

cells.append(nbf.v4.new_code_cell(
"""def soft_threshold(a, t):
    return np.sign(a) * np.maximum(np.abs(a) - t, 0)

train_full = np.loadtxt(f"{DATA}/housing_train.txt")
five = train_full[:5]
RM, MEDV = five[:, 5], five[:, -1]

x = RM - RM.mean()
y5 = MEDV - MEDV.mean()
a1 = x @ y5
norm2 = x @ x
print(f"a^1 = {a1:.3f}   ||x^1||^2 = {norm2:.4f}")

w1_ols = a1 / norm2
w1_ridge = a1 / (norm2 + 1.0)
print(f"w1_OLS   = {w1_ols:.2f}")
print(f"w1_ridge (lambda=1) = {w1_ridge:.2f}")

for lam1 in [5, 10]:
    thresh = lam1 / norm2
    w1_lasso = soft_threshold(a1 / norm2, thresh)
    print(f"lambda_1={lam1:>2}:  threshold={thresh:.2f}   w1_lasso={w1_lasso:.4f}")"""
))

cells.append(nbf.v4.new_markdown_cell(
"""## Part 2 — full coordinate descent on all 13 features, checked against `sklearn`

Cyclic coordinate descent, general (non-unit-norm) update from the note:
$w^j \\leftarrow S_{\\lambda_1/\\|\\mathbf{x}^j\\|^2}(a^j/\\|\\mathbf{x}^j\\|^2)$.

**A real gotcha worth knowing about**: `sklearn`'s `Lasso(alpha=...)` minimizes
$\\frac{1}{2N}\\|\\mathbf{y}-X\\mathbf{w}\\|^2+\\alpha\\|\\mathbf{w}\\|_1$ (mean squared error), while the note's
$J_{\\lambda_1}(w)=\\frac12\\|\\mathbf{y}-X\\mathbf{w}\\|^2+\\lambda_1\\|\\mathbf{w}\\|_1$ uses the *sum*, not the mean. Same
objective shape, different scaling — to compare against `sklearn` apples-to-apples, convert
with $\\alpha = \\lambda_1/N$. Skipping this conversion is a classic silent-mismatch bug: both solvers
run fine, they just silently solve different problems."""
))

cells.append(nbf.v4.new_code_cell(
"""def lasso_coordinate_descent(X, y, lam1, w_init=None, n_iter=200, tol=1e-9):
    N, D = X.shape
    w = np.zeros(D) if w_init is None else w_init.copy()
    col_norm2 = (X ** 2).sum(axis=0)
    for _ in range(n_iter):
        w_old = w.copy()
        for j in range(D):
            r_j = y - X @ w + X[:, j] * w[j]          # partial residual
            a_j = X[:, j] @ r_j
            w[j] = soft_threshold(a_j / col_norm2[j], lam1 / col_norm2[j])
        if np.max(np.abs(w - w_old)) < tol:
            break
    return w

X_train, y_train = train_full[:, :-1], train_full[:, -1]
scaler = StandardScaler().fit(X_train)
X = scaler.transform(X_train)
y = y_train - y_train.mean()   # centered, no intercept needed (see REGRESSION_claude_0926 S6.1)

lam1 = 50.0
w_scratch = lasso_coordinate_descent(X, y, lam1)

sk = Lasso(alpha=lam1 / X.shape[0], fit_intercept=False, max_iter=10000).fit(X, y)

print("scratch :", np.round(w_scratch, 3))
print("sklearn :", np.round(sk.coef_, 3))
print(f"max abs difference: {np.max(np.abs(w_scratch - sk.coef_)):.4f}")
print(f"coefficients driven to exactly 0 (scratch): {np.sum(np.abs(w_scratch) < 1e-8)} / {len(w_scratch)}")"""
))

cells.append(nbf.v4.new_markdown_cell(
"""## Part 3 — the regularization path

Sweep $\\lambda_1$ from the largest value that zeroes every coefficient ($\\lambda_{1,max}=\\max_j|(\\mathbf{x}^j)^T\\mathbf{y}|$,
the smallest correlation any feature would need to survive) down to a small value, warm-starting each
solve from the previous one. Watch features enter one (or a few) at a time — this is the picture
behind "Lasso does automatic feature selection," made concrete instead of asserted."""
))

cells.append(nbf.v4.new_code_cell(
"""feature_names = ["CRIM", "ZN", "INDUS", "CHAS", "NOX", "RM", "AGE", "DIS", "RAD", "TAX", "PTRATIO", "B", "LSTAT"]

lam1_max = np.max(np.abs(X.T @ y))
lam1s = np.linspace(lam1_max, 0.5, 60)

D = X.shape[1]
path = np.zeros((len(lam1s), D))
w = np.zeros(D)
for i, lam1 in enumerate(lam1s):
    w = lasso_coordinate_descent(X, y, lam1, w_init=w)
    path[i] = w

plt.figure(figsize=(7, 5))
for j in range(D):
    plt.plot(lam1s, path[:, j], label=feature_names[j])
plt.gca().invert_xaxis()   # lambda_1 decreasing left-to-right = penalty relaxing
plt.xlabel(r"$\\lambda_1$ (penalty strength, decreasing $\\rightarrow$)")
plt.ylabel("coefficient value")
plt.title("Lasso regularization path (Housing, coordinate descent from scratch)")
plt.axhline(0, color="black", linewidth=0.6)
plt.legend(fontsize=7, ncol=2, loc="lower left")
plt.tight_layout()
plt.show()

nonzero_counts = (np.abs(path) > 1e-6).sum(axis=1)
print("features active at largest lambda_1:", nonzero_counts[0])
print("features active at smallest lambda_1:", nonzero_counts[-1])"""
))

cells.append(nbf.v4.new_markdown_cell(
"""**Reading the plot**: at $\\lambda_1=\\lambda_{1,max}$ every coefficient is exactly $0$ (by construction). As
$\\lambda_1$ relaxes, features cross their individual thresholds and peel off zero one at a time — never
all at once, and not necessarily in a fixed order. `LSTAT` and `RM` (both strong, well-known
predictors of Boston housing prices) are among the first to become nonzero; several others stay at
exactly $0$ for most of the path. This is the same feature-selection story as the note's Section 4.2,
just watched happen continuously instead of at one fixed $\\lambda_1$."""
))

nb["cells"] = cells
nb["metadata"] = {
    "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
    "language_info": {"name": "python", "version": "3.10"},
}

with open("lasso_coordinate_descent.ipynb", "w") as f:
    nbf.write(nb, f)
print("wrote lasso_coordinate_descent.ipynb")