"""
finetuning.py -- numerical companion to the entry 'fine-tuning'.

The entry's claims, measured: a linear hypothesis is pretrained on a large
dataset from a related task, then adapted to a small task-specific training
set by a few GD steps. Warm-starting at the pretrained model
parameters keeps the iterates in a small neighborhood and reaches a small
validation error; the same iteration from a fresh initialization overfits
the small training set. Self-contained (numpy/matplotlib only),
deterministic.

Blocks
------
[B-tasks]    Two related tasks with 30 model parameters: pretraining data
             (2000 data points) from parameters w*, task data from
             w* + delta with a small shift delta. Pretraining recovers
             parameters close to w*.
[B-warm]     Fine-tuning: GD on the task's small training
             set (15 data points), warm-started at the pretrained
             parameters, dips to a validation error below twice the noise
             level within the first ~20 steps; running all 500 steps
             interpolates the 15 points and drifts upward again -- the
             measured reason fine-tuning stops early.
[B-fresh]    The same iteration from a fresh (zero) initialization drives
             the training error low but its validation error stays several
             times larger: it overfits the 15 data points.
[B-distance] The iterates stay near where they start: the distance moved
             is bounded by the accumulated update lengths, and the warm start moves only about the length of the
             task shift delta, far less than the fresh run.

Outputs
-------
pythondemos/finetuning.png : preview figure (checking only).
"""

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from pathlib import Path

OUT_DIR = Path(__file__).parent

report = []


def check(name, ok):
    report.append((name, bool(ok)))
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


D = 30
NOISE = 0.5
gen = np.random.default_rng(20260829)

w_star = gen.normal(size=D)
w_star *= 3.0 / np.linalg.norm(w_star)          # |w*| = 3
delta = gen.normal(size=D)
delta *= 0.5 / np.linalg.norm(delta)            # small task shift, |delta| = 0.5
w_task = w_star + delta


def draw(n, w, seed):
    g = np.random.default_rng(seed)
    X = g.normal(size=(n, D))
    y = X @ w + g.normal(0.0, NOISE, n)
    return X, y


def avg_sqerr(X, y, w):
    return float(np.mean((y - X @ w) ** 2))


# ----------------------------------------------------------------- [B-tasks]
print("[B-tasks] a large pretraining dataset and a small task dataset")

Xp, yp = draw(2000, w_star, seed=1)
w_pre = np.linalg.lstsq(Xp, yp, rcond=None)[0]
Xt, yt = draw(15, w_task, seed=2)
Xv, yv = draw(100_000, w_task, seed=3)
print(f"    pretraining on {len(yp)} data points: |w_pre - w*| = "
      f"{np.linalg.norm(w_pre - w_star):.3f} (|w*| = 3, task shift "
      f"|delta| = 0.5)")
print(f"    task training set: {len(yt)} data points, {D} model parameters")
check("pretraining recovers parameters close to w*",
      np.linalg.norm(w_pre - w_star) < 0.15)


def run_gd(w0, steps=500, eta=0.02):
    w = w0.copy()
    dist, val, moved = [], [], 0.0
    for t in range(steps):
        g = -2.0 / len(yt) * Xt.T @ (yt - Xt @ w)
        w = w - eta * g
        moved += eta * float(np.linalg.norm(g))
        dist.append(float(np.linalg.norm(w - w0)))
        val.append(avg_sqerr(Xv, yv, w))
    return w, dist, val, moved


# ------------------------------------------------------------------ [B-warm]
print("\n[B-warm] warm-started GD on the 15 data points")

w_warm, dist_w, val_w, moved_w = run_gd(w_pre)
best = int(np.argmin(val_w))
print(f"    validation error dips to {val_w[best]:.3f} at step {best + 1} "
      f"(noise level {NOISE**2:.2f}), ends at {val_w[-1]:.3f} after "
      f"interpolating the 15 points (training error "
      f"{avg_sqerr(Xt, yt, w_warm):.3f})")
check("stopped at its best step, the warm start reaches a validation "
      "error below twice the noise level", val_w[best] < 2 * NOISE**2)
check("running on drifts upward again -- the reason fine-tuning is "
      "stopped after few steps", val_w[-1] > 1.2 * val_w[best])


# ----------------------------------------------------------------- [B-fresh]
print("\n[B-fresh] the same iteration from a fresh initialization")

w_fresh, dist_f, val_f, moved_f = run_gd(np.zeros(D))
print(f"    validation error {val_f[-1]:.3f}; training error "
      f"{avg_sqerr(Xt, yt, w_fresh):.3f} (15 data points, 30 parameters)")
check("the fresh run drives the training error below the noise level",
      avg_sqerr(Xt, yt, w_fresh) < NOISE**2)
check("yet its validation error is at least four times the warm start's",
      val_f[-1] > 4 * val_w[-1])


# -------------------------------------------------------------- [B-distance]
print("\n[B-distance] the iterates stay near where they start")

print(f"    distance moved: warm {dist_w[-1]:.2f} (bound "
      f"eta*sum|grad| = {moved_w:.2f}), fresh {dist_f[-1]:.2f}; "
      f"task shift |delta| = 0.5, |w*| = 3")
check("the distance moved is bounded by the accumulated step lengths",
      dist_w[-1] <= moved_w + 1e-9 and dist_f[-1] <= moved_f + 1e-9)
check("the warm start moves on the scale of the task shift, far less "
      "than the fresh run", dist_w[-1] < 1.0 and dist_f[-1] > 2.0)


# --------------------------------------------------------------- preview
fig, ax = plt.subplots(1, 2, figsize=(9.2, 3.6))
steps = range(1, len(val_w) + 1)
ax[0].plot(steps, val_w, "-", color="black", label="warm start (pretrained)")
ax[0].plot(steps, val_f, "--", color="0.4", label="fresh initialization")
ax[0].axhline(NOISE**2, color="0.75", lw=0.8)
ax[0].set_yscale("log")
ax[0].set_xlabel("GD step")
ax[0].set_ylabel("validation error")
ax[0].set_title("[B-warm]/[B-fresh] 15 data points, 30 parameters",
                fontsize=9)
ax[0].legend(frameon=False, fontsize=8)

ax[1].plot(steps, dist_w, "-", color="black", label="warm start")
ax[1].plot(steps, dist_f, "--", color="0.4", label="fresh initialization")
ax[1].set_xlabel("GD step")
ax[1].set_ylabel("distance from initialization")
ax[1].set_title("[B-distance] how far the iterates move", fontsize=9)
ax[1].legend(frameon=False, fontsize=8)

fig.tight_layout()
fig.savefig(OUT_DIR / "finetuning.png", dpi=110)

n_ok = sum(ok for _, ok in report)
print(f"\n{n_ok}/{len(report)} checks pass")
print("wrote finetuning.png")
if n_ok != len(report):
    raise SystemExit(1)