"""
rv.py — numerical companion to the glossary entry 'random variable (RV)'.

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]   An RV is a function on the sample space of a random
          experiment: for a fair-die experiment (sample space
          {1,...,6}) the RV x(omega) = 1{omega is even} is an explicit
          function; simulating the experiment and applying the function
          reproduces the probability P(x = 1) = 1/2 implied by the
          uniform distribution on the sample space.
[P-types] The types listed in the entry — binary RV, discrete RV,
          real-valued RV, random vector, random matrix — are
          instantiated as functions of the same underlying experiment,
          with values in {0, 1}, a countable set, R, R^d, and
          R^{m x d}, respectively.

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

Data generated by pythondemos/rv.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] an RV is a function on the sample space")
sample_space = np.arange(1, 7)                   # fair die
x_of = lambda omega: (omega % 2 == 0).astype(int)  # RV: even -> 1
outcomes = rng.choice(sample_space, size=10**6)  # run the experiment
x_real = x_of(outcomes)                          # realizations of the RV
check("x is a deterministic function of the outcome",
      np.array_equal(x_real, x_of(outcomes)))
check("P(x = 1) = 1/2 from the uniform experiment",
      abs(x_real.mean() - 0.5) < 2e-3)
check("the same outcome always maps to the same value",
      x_of(np.array([4]))[0] == 1 and x_of(np.array([3]))[0] == 0)

# --------------------------------------------------------- [P-types]
print("[P-types] binary / discrete / real-valued / vector / matrix RVs")
omega = rng.uniform(size=10**4)                  # one underlying experiment
binary = (omega > 0.5).astype(int)
discrete = np.floor(10 * omega).astype(int)      # values in {0,...,9}
realval = -np.log(omega)                         # values in R (exponential)
vec = np.stack([omega, omega**2, np.sin(omega)], axis=1)      # R^3
mat = omega[:, None, None] * np.ones((1, 2, 3))               # R^{2x3}
check("binary RV takes values in {0, 1}",
      set(np.unique(binary)) <= {0, 1})
check("discrete RV takes values in a countable set",
      np.issubdtype(discrete.dtype, np.integer)
      and len(np.unique(discrete)) <= 10)
check("real-valued RV takes values in R (nonnegative here)",
      realval.dtype == float and np.all(realval >= 0))
check("random vector maps outcomes to R^3", vec.shape == (10**4, 3))
check("random matrix maps outcomes to R^{2x3}",
      mat.shape == (10**4, 2, 3))

# ------------------------------------------------------------ preview
# The outcomes and the RV values live on different sets, so they get one
# panel each: drawing them on shared x positions would hide one series.
fig, ax = plt.subplots(1, 2, figsize=(7.0, 3.0))
vals, counts = np.unique(outcomes[:600], return_counts=True)
ax[0].bar(vals, counts / 600, width=0.6, color="0.75", edgecolor="black")
ax[0].set_xticks(range(1, 7))
ax[0].set_xlabel("outcome of the die roll")
ax[0].set_ylabel("relative frequency")
ax[0].set_title("[P-def] the 6 outcomes")
ax[1].bar([0, 1], [np.mean(x_real == 0), np.mean(x_real == 1)],
          width=0.6, color="white", edgecolor="black", hatch="///")
ax[1].set_xticks([0, 1])
ax[1].set_xlabel("value of the RV 1{even}")
ax[1].set_ylabel("relative frequency")
ax[1].set_title("[P-def] the 2 values it takes")
fig.tight_layout()
fig.savefig("rv.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)