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

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-def]    The variance of a real-valued RV is E{(x - E{x})^2}: the
           empirical variance of iid Gaussian draws matches the analytic
           sigma^2, and shifting the RV leaves the variance unchanged
           (it measures spread around the mean, not location).
[P-vector] For a random vector, E{||x - E{x}||^2} equals trace(C), the
           sum of the per-entry variances (checked against the analytic
           covariance matrix C = A A^T and against np.cov).

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

Data generated by pythondemos/variance.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}")


# ----------------------------------------------------------- [P-def]
print("[P-def] variance = E{(x - E x)^2}, invariant under shifts")
sigma = 1.7
m = 10**6
x = sigma * rng.standard_normal(m)
var_emp = np.mean((x - x.mean()) ** 2)
check("empirical variance matches sigma^2 = 2.89",
      abs(var_emp - sigma**2) < 2e-2)
check("shifting by +5 leaves the variance unchanged",
      abs(np.mean(((x + 5) - (x + 5).mean()) ** 2) - var_emp) < 1e-9)
scales = [0.5, 1.0, 2.0]
vars_scaled = [np.mean(((s * x) - (s * x).mean()) ** 2) for s in scales]
check("scaling by s multiplies the variance by s^2",
      all(abs(v - s**2 * var_emp) < 1e-6 * max(1, s**2 * var_emp)
          for s, v in zip(scales, vars_scaled)))

# -------------------------------------------------------- [P-vector]
print("[P-vector] E{||x - E x||^2} = trace(C) = sum of entry variances")
A = np.array([[1.0, 0.0, 0.0], [0.5, 0.8, 0.0], [-0.2, 0.3, 1.1]])
C = A @ A.T                                   # analytic covariance
z = rng.standard_normal((10**6, 3))
xv = z @ A.T                                  # zero-mean, covariance C
sq_dev = np.mean(np.sum((xv - xv.mean(axis=0)) ** 2, axis=1))
check("empirical E||x - Ex||^2 matches trace(C)",
      abs(sq_dev - np.trace(C)) < 2e-2)
per_entry = np.array([np.mean((xv[:, j] - xv[:, j].mean()) ** 2)
                      for j in range(3)])
check("sum of per-entry variances equals trace(C)",
      abs(per_entry.sum() - np.trace(C)) < 2e-2)
check("per-entry variances match diag(np.cov)",
      np.allclose(per_entry, np.diag(np.cov(xv.T, ddof=0)), atol=1e-9))

# ------------------------------------------------------------ preview
fig, ax = plt.subplots(figsize=(4.5, 3.2))
ax.bar(range(3), per_entry, label="per-entry variance")
ax.axhline(np.trace(C), ls="--", c="k",
           label=f"trace(C) = {np.trace(C):.2f}")
ax.plot([0, 1, 2], np.cumsum(per_entry), "o-", c="C1",
        label="cumulative sum")
ax.set_xlabel("entry j"); ax.legend(frameon=False)
ax.set_title("[P-vector] variances sum to trace(C)")
fig.tight_layout()
fig.savefig("variance.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)