"""
covmtx.py — numerical companion to the glossary entry 'covariance matrix'.

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 covariance matrix C = E{(x - E x)(x - E x)^T}: the empirical
        outer-product average of iid draws from a random vector with
        analytic covariance A A^T recovers that matrix; its entry (j, j')
        is the covariance of entries j and j' (checked against np.cov),
        and its diagonal holds the per-entry variances.

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

Data generated by pythondemos/covmtx.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] C = E{(x - Ex)(x - Ex)^T} — empirical vs analytic")
A = np.array([[1.0, 0.0, 0.0], [0.5, 0.8, 0.0], [-0.2, 0.3, 1.1]])
mu = np.array([2.0, -1.0, 0.5])
C = A @ A.T                                    # analytic covariance
m = 10**6
x = rng.standard_normal((m, 3)) @ A.T + mu
xc = x - x.mean(axis=0)                        # centered
C_emp = (xc[:, :, None] * xc[:, None, :]).mean(axis=0)
check("empirical outer-product average recovers C (|err| < 5e-3)",
      np.max(np.abs(C_emp - C)) < 5e-3)
check("matches np.cov (ddof=0)",
      np.allclose(C_emp, np.cov(x.T, ddof=0), atol=1e-9))
covs = np.array([[np.mean(xc[:, j] * xc[:, k]) for k in range(3)]
                 for j in range(3)])
check("entry (j, j') is the covariance of entries j and j'",
      np.allclose(covs, C_emp, atol=1e-12))
check("diagonal entries are the per-entry variances",
      np.allclose(np.diag(C_emp),
                  [np.mean(xc[:, j] ** 2) for j in range(3)], atol=1e-12))

# ------------------------------------------------------------ preview
fig, ax = plt.subplots(1, 2, figsize=(7.5, 3.0))
for a, M, t in ((ax[0], C, "analytic C = A A^T"),
                (ax[1], C_emp, "empirical (m = 1e6)")):
    im = a.imshow(M, cmap="gray")
    a.set_title(t)
    fig.colorbar(im, ax=a, shrink=0.8)
fig.suptitle("[P-def] covariance matrix")
fig.tight_layout()
fig.savefig("covmtx.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)