"""
norm.py — numerical companion to the glossary entry 'norm'.

One block per paragraph of the entry (marked [P...]): each block verifies
numerically what the corresponding statement asserts. Self-contained
(numpy/matplotlib only), fixed seed.

Blocks
------
[P-axioms] The three norm axioms — definiteness, homogeneity, triangle
           inequality — hold for the l1-, l2-, and linf-norms on random
           vector pairs (10^4 trials each).
[P-metric] d(u, v) = ||u - v|| is a metric: symmetry, identity of
           indiscernibles, and the triangle inequality checked on random
           triples.
[P-inner]  The inner product induces the l2-norm: ||u|| = sqrt(<u, u>)
           for random vectors.
[P-lp]     The lp-norm family: the explicit sum formula matches
           np.linalg.norm for p in {1, 2, inf}, the linf-norm is the
           max absolute entry, and the unit spheres of l1/l2/linf are
           nested (||x||_inf <= ||x||_2 <= ||x||_1) — the geometry of
           the entry's unit-ball figure.
[P-ml]     Norms define losses and regularizers: the squared-error loss
           is a squared l2-norm of the residual, and the l1/l2
           regularizers evaluate norms of
           the parameter vector.

Outputs
-------
norm.png : preview figure (checking only).

Data generated by pythondemos/norm.py.
"""

import numpy as np
import matplotlib

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

rng = np.random.default_rng(42)
report = []


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


NORMS = {"l1": 1, "l2": 2, "linf": np.inf}
U = rng.normal(size=(10**4, 5))
V = rng.normal(size=(10**4, 5))

# -------------------------------------------------------- [P-axioms]
print("[P-axioms] definiteness, homogeneity, triangle inequality")
for name, p in NORMS.items():
    nu = np.linalg.norm(U, p, axis=1)
    check(f"{name}: norm(0) = 0 and norm(u) > 0 for u != 0",
          np.linalg.norm(np.zeros(5), p) == 0 and np.all(nu > 0))
    a = rng.normal()
    check(f"{name}: homogeneity ||a u|| = |a| ||u||",
          np.allclose(np.linalg.norm(a * U, p, axis=1), abs(a) * nu))
    check(f"{name}: triangle ||u + v|| <= ||u|| + ||v||",
          np.all(np.linalg.norm(U + V, p, axis=1)
                 <= nu + np.linalg.norm(V, p, axis=1) + 1e-12))

# -------------------------------------------------------- [P-metric]
print("[P-metric] d(u, v) = ||u - v|| is a metric")
W = rng.normal(size=(10**4, 5))
d = lambda X, Y: np.linalg.norm(X - Y, 2, axis=1)
check("symmetry d(u, v) = d(v, u)", np.allclose(d(U, V), d(V, U)))
check("d(u, u) = 0", np.all(d(U, U) == 0))
check("d(u, v) > 0 whenever u != v (identity of indiscernibles)",
      np.all(d(U, V)[np.any(U != V, axis=1)] > 0))
check("triangle d(u, w) <= d(u, v) + d(v, w)",
      np.all(d(U, W) <= d(U, V) + d(V, W) + 1e-12))

# --------------------------------------------------------- [P-inner]
print("[P-inner] the inner product induces the l2-norm")
check("||u|| = sqrt(<u, u>)",
      np.allclose(np.linalg.norm(U, 2, axis=1),
                  np.sqrt(np.sum(U * U, axis=1))))

# ------------------------------------------------------------ [P-lp]
print("[P-lp] the lp family and its unit-ball geometry")
x = rng.normal(size=(10**4, 5))
lp_sum = lambda X, p: (np.sum(np.abs(X) ** p, axis=1)) ** (1 / p)
check("sum formula matches np.linalg.norm for p = 1, 2",
      np.allclose(lp_sum(x, 1), np.linalg.norm(x, 1, axis=1))
      and np.allclose(lp_sum(x, 2), np.linalg.norm(x, 2, axis=1)))
check("linf-norm is the max absolute entry",
      np.allclose(np.linalg.norm(x, np.inf, axis=1),
                  np.max(np.abs(x), axis=1)))
check("||x||_inf <= ||x||_2 <= ||x||_1 (nested unit balls)",
      np.all(np.linalg.norm(x, np.inf, axis=1)
             <= np.linalg.norm(x, 2, axis=1) + 1e-12)
      and np.all(np.linalg.norm(x, 2, axis=1)
                 <= np.linalg.norm(x, 1, axis=1) + 1e-12))

# ------------------------------------------------------------ [P-ml]
print("[P-ml] norms define losses and regularizers")
Xf = rng.normal(size=(50, 3))
yf = Xf @ np.array([1.0, 0.0, -0.5]) + 0.1 * rng.normal(size=50)
w = np.linalg.lstsq(Xf, yf, rcond=None)[0]
sq_loss = np.mean((yf - Xf @ w) ** 2)
check("squared-error loss = (1/m) ||y - X w||_2^2",
      np.isclose(sq_loss, np.linalg.norm(yf - Xf @ w) ** 2 / 50))
check("ridge regularizer alpha ||w||_2^2 and Lasso regularizer "
      "alpha ||w||_1 are norm evaluations",
      np.isclose(np.linalg.norm(w, 2) ** 2, np.sum(w**2))
      and np.isclose(np.linalg.norm(w, 1), np.sum(np.abs(w))))

# ------------------------------------------------------------ preview
th = np.linspace(0, 2 * np.pi, 400)
circ = np.stack([np.cos(th), np.sin(th)])
fig, ax = plt.subplots(figsize=(4.0, 4.0))
ax.plot(*(circ / np.linalg.norm(circ, 1, axis=0)), label="$\\ell_1$")
ax.plot(*circ, label="$\\ell_2$")
ax.plot(*(circ / np.linalg.norm(circ, np.inf, axis=0)),
        label="$\\ell_\\infty$")
ax.set_aspect("equal"); ax.legend(frameon=False)
ax.set_title("[P-lp] unit spheres")
fig.tight_layout()
fig.savefig("norm.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)