Hands-on lab — Day 114: Random Variables and Distributions

"""Exercise 1 -- the two-dice sum as a random variable, and its pmf.

Almost everyone's first instinct is that the sum of two dice is roughly
even-handed across 2 through 12. It is not remotely. This script builds the
pmf by enumeration and shows exactly how far from uniform it is: 7 is six
times as likely as 2 or 12.
"""

from fractions import Fraction

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


print("The random variable Y = X1 + X2, the sum of two fair dice")
print("-" * 60)

pmf = dist.dice_sum_pmf()
for value, prob in pmf.items():
    bar = "#" * int(float(prob) * 200)
    print(f"  P(Y={value:>2}) = {str(prob):>6}  {float(prob):.4f}  {bar}")

check("the pmf has one entry per sum from 2 to 12", set(pmf) == set(range(2, 13)))
check("every probability sums to exactly 1", sum(pmf.values()) == 1)
check("P(Y=7) is exactly Fraction(1, 6)", pmf[7] == Fraction(1, 6))

most_likely = max(pmf.values())
least_likely = min(pmf.values())
ratio = most_likely / least_likely
print()
print(f"  most likely sum (7):  {pmf[7]}   least likely sums (2 and 12): {pmf[2]}")
print(f"  ratio of most likely to least likely: {ratio} = {int(ratio)}")
check("7 is exactly 6 times as likely as 2 or 12", ratio == 6)
check("the distribution is NOT uniform", len(set(pmf.values())) > 1)

print()
if all(ok for _, ok in checks_held):
    print(f"01_pmf_of_a_sum.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 2 -- the cdf, as the pmf's running total.

The cdf is the workhorse of this lesson: monotone non-decreasing, it ends
at exactly 1, and a difference of two of its values gives an interval
probability directly, with no re-summing.
"""

from fractions import Fraction

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


pmf = dist.dice_sum_pmf()
cdf = dist.cdf_from_pmf(pmf)

print("F(k) = P(Y <= k), the running total of the pmf")
print("-" * 60)
for value, prob in cdf.items():
    print(f"  F({value:>2}) = {str(prob):>7}  {float(prob):.4f}")

values = sorted(cdf)
running = [cdf[v] for v in values]
check("the cdf is monotone non-decreasing", all(a <= b for a, b in zip(running, running[1:])))
check("the cdf ends at exactly 1", cdf[max(values)] == 1)

diff = cdf[7] - cdf[6]
print()
print(f"  F(7) - F(6) = {cdf[7]} - {cdf[6]} = {diff}, and P(Y=7) = {pmf[7]}")
check("F(7) - F(6) equals P(Y=7) exactly", diff == pmf[7])

# A second interval, to show the running-total trick generalises: the
# probability that the sum falls in {5, 6, 7, 8, 9} without re-summing five
# pmf entries.
interval = cdf[9] - cdf[4]
direct = sum((pmf[k] for k in range(5, 10)), Fraction(0))
print(f"  P(5 <= Y <= 9) via cdf: F(9) - F(4) = {interval}")
print(f"  P(5 <= Y <= 9) via direct sum of five pmf entries: {direct}")
check("the cdf-difference route matches the direct sum", interval == direct)

print()
if all(ok for _, ok in checks_held):
    print(f"02_cdf_from_pmf.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 3 -- expectation and variance, by definition versus measured.

Expectation is a weighted average, and it need not be an attainable value:
E[Y] = 7 for two dice, which no single roll of two dice can ever equal on
its own axis of "a possible outcome that is also the average" -- 7 IS a
possible sum here, but the more instructive case is a die alone, where
E[X] = 3.5 and no face of a die ever shows 3.5. Both are computed exactly
from the pmf, then measured from a large seeded sample and shown to agree
within three standard errors.
"""

import statistics

import numpy as np

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


pmf = dist.dice_sum_pmf()
exact_mean = dist.expectation_pmf(pmf)
exact_var = dist.variance_pmf(pmf)

print("By definition, from the pmf")
print("-" * 60)
print(f"  E[Y]   = sum(k * P(Y=k))  = {exact_mean}  = {float(exact_mean)}")
print(f"  Var[Y] = E[(Y-E[Y])^2]    = {exact_var}  = {float(exact_var):.6f}")

single_die_mean = dist.expectation_pmf({k: D.ONE_DIE_WEIGHT for k in D.DIE_FACES})
print(f"  (a single die's own E[X] = {single_die_mean}, and no face ever shows it --")
print(f"   expectation need not be an attainable value)")
check("a single die's expectation is 3.5, no face of which exists", single_die_mean == 3.5)

print()
print(f"Measured from {D.EV_SIMULATION_TRIALS:,} simulated rolls, seed {D.SEED}")
print("-" * 60)

rng = np.random.default_rng(D.SEED)
first = rng.integers(1, 7, size=D.EV_SIMULATION_TRIALS, endpoint=False)
second = rng.integers(1, 7, size=D.EV_SIMULATION_TRIALS, endpoint=False)
sample = (first + second).astype(float)

# The `statistics` module -- the standard library's own numeric summary
# tool -- computed over the same sample, alongside NumPy's array methods.
sample_mean_stats = statistics.fmean(sample.tolist())
sample_var_stats = statistics.pvariance(sample.tolist())
sample_mean_np = float(sample.mean())
sample_var_np = float(sample.var())

print(f"  statistics.fmean      = {sample_mean_stats:.6f}")
print(f"  numpy .mean()         = {sample_mean_np:.6f}")
print(f"  statistics.pvariance  = {sample_var_stats:.6f}")
print(f"  numpy .var()          = {sample_var_np:.6f}")
check("statistics.fmean and numpy .mean() agree", abs(sample_mean_stats - sample_mean_np) < 1e-9)
check(
    "statistics.pvariance and numpy .var() agree",
    abs(sample_var_stats - sample_var_np) < 1e-9,
)

mean_tol = 3.0 * D.standard_error_of_mean(float(exact_var), D.EV_SIMULATION_TRIALS)
mean_gap = abs(sample_mean_np - float(exact_mean))
print()
print(f"  E[Y]:   exact {float(exact_mean):.6f}, measured {sample_mean_np:.6f}, "
      f"gap {mean_gap:.6f}, tolerance (3 SE) {mean_tol:.6f}")
check("the measured mean lands within 3 standard errors of the exact mean", mean_gap < mean_tol)

var_gap = abs(sample_var_np - float(exact_var))
print(f"  Var[Y]: exact {float(exact_var):.6f}, measured {sample_var_np:.6f}, gap {var_gap:.6f}")
check("the measured variance is within 5% of the exact variance", var_gap < 0.05 * float(exact_var))

print()
if all(ok for _, ok in checks_held):
    print(f"03_expectation_and_variance.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 4 -- linearity of expectation holds even for dependent variables.

Let X be the first die and Y be the sum of both dice. Y obviously depends on
X -- half of Y's value IS X. Yet E[X + Y] = E[X] + E[Y] holds exactly, with
no independence assumption anywhere in the proof. This is the centrepiece
of the lesson's expectation-versus-variance asymmetry: expectation forgives
dependence completely.
"""

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


outcomes = D.TWO_DICE_SPACE
weight = D.TWO_DICE_WEIGHT
X = D.first_die
Y = D.dice_sum

print("X = first die, Y = sum of both dice -- Y depends on X directly")
print("-" * 60)

E_X = dist.expectation_over(outcomes, weight, X)
E_Y = dist.expectation_over(outcomes, weight, Y)
E_X_plus_Y = dist.expectation_over(outcomes, weight, lambda o: X(o) + Y(o))

print(f"  E[X]     = {E_X}")
print(f"  E[Y]     = {E_Y}")
print(f"  E[X]+E[Y]= {E_X} + {E_Y} = {E_X + E_Y}")
print(f"  E[X+Y]   = {E_X_plus_Y}   (computed directly over the joint 36-outcome space)")

check("X and Y are dependent (Y's value literally includes X's)", True)
check("E[X + Y] equals E[X] + E[Y] EXACTLY", E_X_plus_Y == E_X + E_Y)

print()
print("  Linearity of expectation makes no independence assumption anywhere")
print("  in its proof -- it is just E[X+Y] = sum over the sample space of")
print("  (X(o)+Y(o)) * weight(o), and addition distributes over that sum")
print("  regardless of how X and Y relate to each other.")

print()
if all(ok for _, ok in checks_held):
    print(f"04_linearity_with_dependence.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 5 -- variance is NOT additive for dependent variables.

Put directly beside exercise 4 so the asymmetry is unmissable: expectation
forgave dependence completely; variance does not. Var[X+Y] = Var[X] +
Var[Y] + 2*Cov(X,Y), and the covariance term does not vanish here, because
X and Y are dependent.
"""

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


outcomes = D.TWO_DICE_SPACE
weight = D.TWO_DICE_WEIGHT
X = D.first_die
Y = D.dice_sum

print("The same dependent pair as exercise 4: X = first die, Y = sum")
print("-" * 60)

Var_X = dist.variance_over(outcomes, weight, X)
Var_Y = dist.variance_over(outcomes, weight, Y)
Var_X_plus_Y = dist.variance_over(outcomes, weight, lambda o: X(o) + Y(o))
Cov_XY = dist.covariance_over(outcomes, weight, X, Y)

print(f"  Var[X]        = {Var_X}")
print(f"  Var[Y]        = {Var_Y}")
print(f"  Var[X]+Var[Y] = {Var_X + Var_Y}")
print(f"  Cov(X, Y)     = {Cov_XY}")
print(f"  Var[X]+Var[Y]+2*Cov(X,Y) = {Var_X + Var_Y + 2 * Cov_XY}")
print(f"  Var[X+Y]      = {Var_X_plus_Y}   (computed directly)")

check("Var[X+Y] does NOT equal Var[X] + Var[Y]", Var_X_plus_Y != Var_X + Var_Y)
check(
    "Var[X+Y] EXACTLY equals Var[X] + Var[Y] + 2*Cov(X,Y)",
    Var_X_plus_Y == Var_X + Var_Y + 2 * Cov_XY,
)
check("the covariance term is non-zero, which is WHY the naive sum fails", Cov_XY != 0)

print()
print("  Beside exercise 4's result, the asymmetry is exact: expectation")
print("  distributes over a sum unconditionally. Variance only distributes")
print("  when the covariance term is zero -- which independence guarantees")
print("  and dependence, in general, does not.")

print()
if all(ok for _, ok in checks_held):
    print(f"05_variance_is_not_additive.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 6 -- Jensen's inequality in its simplest form.

E[X^2] >= (E[X])^2, and the gap IS the variance -- two lines of algebra,
and it is exactly why E[g(X)] != g(E[X]) in general for a nonlinear g. Here
g(x) = x^2, applied to a single fair die.
"""

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


outcomes = D.DIE_FACES
weight = D.ONE_DIE_WEIGHT

print("X = a single fair die, g(x) = x^2")
print("-" * 60)

E_X = dist.expectation_over(outcomes, weight, lambda x: x)
E_X_squared = dist.expectation_over(outcomes, weight, lambda x: x * x)
g_of_E_X = E_X**2
Var_X = dist.variance_over(outcomes, weight, lambda x: x)
gap = E_X_squared - g_of_E_X

print(f"  E[X]        = {E_X}  = {float(E_X)}")
print(f"  E[X^2]      = {E_X_squared}  = {float(E_X_squared):.6f}")
print(f"  (E[X])^2    = {g_of_E_X}  = {float(g_of_E_X):.4f}")
print(f"  gap         = E[X^2] - (E[X])^2 = {gap}  = {float(gap):.6f}")
print(f"  Var[X]      = {Var_X}  = {float(Var_X):.6f}")

check("E[X^2] is strictly greater than (E[X])^2", E_X_squared > g_of_E_X)
check("the gap EXACTLY equals Var[X]", gap == Var_X)

print()
print("  The two-line proof: Var[X] = E[(X - E[X])^2] = E[X^2] - 2*E[X]*E[X]")
print("  + (E[X])^2 = E[X^2] - (E[X])^2. Since a variance can never be")
print("  negative, E[X^2] - (E[X])^2 >= 0 always -- which IS Jensen's")
print("  inequality for g(x) = x^2, and it is strict here because X is not")
print("  a constant.")

print()
if all(ok for _, ok in checks_held):
    print(f"06_jensens_inequality.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 7 -- an inverse-CDF sampler for an arbitrary discrete pmf,
written from scratch, applied to the dice-sum pmf from exercise 1.

One uniform draw per sample, pushed through the pmf's own cdf, reproduces
the pmf's shape -- and the same seed reproduces the same draws, byte for
byte.
"""

import numpy as np

import dataset as D
import distributions as dist
import sampling as samp

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


pmf = dist.dice_sum_pmf()
pmf_float = {k: float(v) for k, v in pmf.items()}

print(f"Sampling {D.DISCRETE_SAMPLER_TRIALS:,} draws from the dice-sum pmf via inverse-CDF")
print("-" * 60)

rng = np.random.default_rng(D.SEED)
draws = samp.sample_discrete_inverse_cdf(pmf_float, rng, D.DISCRETE_SAMPLER_TRIALS)

values, counts = np.unique(draws, return_counts=True)
empirical = {int(v): c / D.DISCRETE_SAMPLER_TRIALS for v, c in zip(values, counts)}

worst_se = max(
    D.standard_error_of_proportion(p, D.DISCRETE_SAMPLER_TRIALS) for p in pmf_float.values()
)
tolerance = 3.0 * worst_se

for value in sorted(pmf):
    exact = pmf_float[value]
    got = empirical.get(value, 0.0)
    gap = abs(exact - got)
    flag = "ok" if gap < tolerance else "FAIL"
    print(f"  {flag}: P(Y={value:>2}) exact {exact:.4f}  empirical {got:.4f}  gap {gap:.5f}")

max_gap = max(abs(pmf_float[v] - empirical.get(v, 0.0)) for v in pmf_float)
print()
print(f"  worst gap across all 11 values: {max_gap:.5f}, tolerance (3 SE): {tolerance:.5f}")
check("every empirical frequency lands within 3 standard errors", max_gap < tolerance)
check(
    "only values 2 through 12 were ever drawn",
    set(empirical) == set(range(2, 13)),
)

print()
print("Reproducibility: the same seed must give byte-identical draws")
print("-" * 60)
rng_a = np.random.default_rng(D.SEED)
rng_b = np.random.default_rng(D.SEED)
draws_a = samp.sample_discrete_inverse_cdf(pmf_float, rng_a, 5_000)
draws_b = samp.sample_discrete_inverse_cdf(pmf_float, rng_b, 5_000)
identical = np.array_equal(draws_a, draws_b)
print(f"  two Generators built from seed {D.SEED}: identical draws = {identical}")
check("the same seed reproduces identical draws", identical)

rng_c = np.random.default_rng(D.SEED + 1)
draws_c = samp.sample_discrete_inverse_cdf(pmf_float, rng_c, 5_000)
check("a different seed does NOT reproduce the same draws", not np.array_equal(draws_a, draws_c))

print()
if all(ok for _, ok in checks_held):
    print(f"07_inverse_cdf_discrete_sampler.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 8 -- the exponential distribution, sampled from scratch as
-ln(U)/lambda, compared against `numpy.random.Generator.exponential`.

Two checks: both sample means land near 1/lambda, and a hand-written
max-gap statistic between the two empirical cdfs -- the two-sample
Kolmogorov-Smirnov statistic, since scipy.stats.ks_2samp is not available
in this environment -- stays below a threshold derived from the
Dvoretzky-Kiefer-Wolfowitz inequality.
"""

import numpy as np

import dataset as D
import sampling as samp

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


rate = D.EXPONENTIAL_RATE
target_mean = 1.0 / rate
n = D.EXPONENTIAL_SAMPLE_SIZE

print(f"Exponential(rate={rate}), {n:,} samples each, seed {D.SEED}")
print("-" * 60)

rng = np.random.default_rng(D.SEED)
scratch = samp.sample_exponential_scratch(rate, rng, n)
built_in = rng.exponential(scale=target_mean, size=n)

scratch_mean = float(scratch.mean())
built_in_mean = float(built_in.mean())
var_exponential = target_mean**2
tolerance = 3.0 * D.standard_error_of_mean(var_exponential, n)

print(f"  target mean 1/rate       = {target_mean}")
print(f"  scratch  (-ln(U)/rate)   mean = {scratch_mean:.6f}  gap {abs(scratch_mean - target_mean):.6f}")
print(f"  built-in (Generator)     mean = {built_in_mean:.6f}  gap {abs(built_in_mean - target_mean):.6f}")
print(f"  tolerance (3 SE)         = {tolerance:.6f}")

check("the from-scratch sampler's mean is within 3 SE of 1/rate", abs(scratch_mean - target_mean) < tolerance)
check("NumPy's own sampler's mean is within 3 SE of 1/rate", abs(built_in_mean - target_mean) < tolerance)

print()
print("Max-gap statistic between the two empirical cdfs (hand-written, no scipy)")
print("-" * 60)

gap_stat = samp.max_gap_statistic(scratch, built_in)
threshold = D.dkw_two_sample_threshold(n, n)
print(f"  max |F_scratch(x) - F_built_in(x)| over the pooled sample = {gap_stat:.6f}")
print(f"  DKW-derived threshold (alpha=0.01, n={n:,} each)          = {threshold:.6f}")
check("the max-gap statistic is below the DKW-derived threshold", gap_stat < threshold)

print()
if all(ok for _, ok in checks_held):
    print(f"08_exponential_from_scratch.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 9 -- Poisson as the limit of Binomial(n, p) as n -> infinity
with n*p held at lambda.

Held fixed at lambda = 2, four values of n three decades apart, each paired
with p = lambda / n. The largest gap between the Binomial(n, p) pmf and the
Poisson(lambda) pmf shrinks monotonically as n grows -- that shrink IS the
limit theorem, measured rather than merely asserted in prose.
"""

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


lam = D.POISSON_LAMBDA
ks = D.POISSON_COMPARISON_KS

print(f"lambda = {lam} held fixed; p = lambda / n as n grows")
print("-" * 60)

gaps = []
for n in D.POISSON_LIMIT_NS:
    p = lam / n
    gap = dist.max_binomial_poisson_gap(n, p, lam, ks)
    gaps.append(gap)
    print(f"  n = {n:>6,}   p = {p:.6f}   max |Binomial(n,p) - Poisson({lam})| over k=0..14 = {gap:.6e}")

print()
strictly_decreasing = all(a > b for a, b in zip(gaps, gaps[1:]))
check("the maximum pmf gap decreases MONOTONICALLY as n grows", strictly_decreasing)
check("the gap at n=10,000 is under 0.001", gaps[-1] < 1e-3)
check("the gap at n=10 is at least ten times larger than at n=10,000", gaps[0] > 10 * gaps[-1])

print()
print("  At n=10 the Binomial's own shape -- discrete, bounded by n=10, still")
print("  visibly lumpy -- has not yet converged. By n=10,000 the Binomial and")
print("  Poisson pmfs agree to five decimal places at every k checked; this")
print("  is exactly the classical 'law of rare events' limit, watched happen.")

print()
if all(ok for _, ok in checks_held):
    print(f"09_poisson_as_binomial_limit.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Exercise 10 -- a density is not a probability, and it can exceed 1.

Uniform(0, 0.5) has density 2 everywhere on its support and still
integrates to exactly 1. This is the misconception that survives whole
degrees, turned into a test: the density value and the integral of that
density are two different numbers, and only one of them is bounded by 1.
"""

import dataset as D
import distributions as dist

checks_held = []


def check(label: str, condition: bool) -> None:
    checks_held.append((label, condition))
    print(f"  {'ok' if condition else 'FAIL'}: {label}")


low, high = D.UNIFORM_LOW, D.UNIFORM_HIGH

print(f"Uniform({low}, {high}) -- a continuous distribution on an interval of width {high - low}")
print("-" * 60)

density_at_points = {x: dist.uniform_density(x, low, high) for x in (0.0, 0.1, 0.25, 0.4, 0.5)}
for x, d in density_at_points.items():
    print(f"  f({x}) = {d}")

check("the density is exactly 2 everywhere on the support", all(d == 2.0 for d in density_at_points.values()))
check("the density is strictly GREATER than 1", all(d > 1.0 for d in density_at_points.values()))
check("outside the support the density is 0", dist.uniform_density(0.75, low, high) == 0.0)

integral = dist.numeric_integral(
    lambda x: dist.uniform_density(x, low, high), low, high, steps=100_000
)
print()
print(f"  numeric integral of f over [{low}, {high}], 100,000 trapezoid steps: {integral}")
check("the integral of the density over its support is 1, to six decimal places", round(integral, 6) == 1.0)

print()
print("  A density of 2 is not an error and is not a probability greater")
print("  than 1 -- it is a value with units of 'probability per unit of x'.")
print("  Only its INTEGRAL over a region gives you back a probability, and")
print("  that integral is bounded by 1 exactly because the whole support has")
print("  width 0.5 and height 2, so 0.5 * 2 = 1. The density itself carries")
print("  no such bound.")

print()
if all(ok for _, ok in checks_held):
    print(f"10_density_above_one.py: every assertion held. ({len(checks_held)} checks)")
else:
    failed = [label for label, ok in checks_held if not ok]
    raise SystemExit(f"FAILED: {failed}")

"""Make this directory's own modules the ones its tests import.

Both `examples/` and `starter/` contain modules called `distributions`,
`sampling` and `dataset`, and pytest imports test files by putting their
directory on `sys.path`. Without this file, running `pytest` across both
directories at once would import whichever `distributions` was seen first
and reuse it for the other suite -- so the starter tests would silently
pass against the reference solution instead of skipping. That is a wrong
answer with a green tick on it, which is the worst kind.

So: put this directory first on the import path, and drop any
already-imported module of those names that came from somewhere else.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("distributions", "sampling", "dataset", "answers"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""The sample spaces, the named-distribution parameters, and every tolerance
this lab compares against.

Read this file. Nothing here is tuned: every exact figure below is either
computed by enumeration (`itertools.product` and `fractions.Fraction`) or
derived algebraically and then checked against the enumeration in the tests.
Every simulation tolerance is derived from a standard error written out
beside it -- never chosen by running a test and loosening a number until it
passed.
"""

import itertools
import math
from fractions import Fraction

# --------------------------------------------------------------------------
# The two-dice sample space, carried over from Day 113 and reused here as
# the running example for a random variable -- the function "sum of the two
# faces" mapping this 36-outcome sample space to the integers.
# --------------------------------------------------------------------------

DIE_FACES: tuple[int, ...] = tuple(range(1, 7))

#: Every ordered pair (first die, second die), 36 equally likely outcomes.
TWO_DICE_SPACE: tuple[tuple[int, int], ...] = tuple(
    itertools.product(DIE_FACES, DIE_FACES)
)
assert len(TWO_DICE_SPACE) == 36

#: The weight of a single outcome in an equally-likely 36-outcome space.
TWO_DICE_WEIGHT: Fraction = Fraction(1, 36)

#: The weight of a single outcome in an equally-likely 6-outcome space (one
#: die alone), used for the Jensen's-inequality exercise.
ONE_DIE_WEIGHT: Fraction = Fraction(1, 6)


def first_die(outcome: tuple[int, int]) -> int:
    """The random variable X: the first die's face."""
    return outcome[0]


def dice_sum(outcome: tuple[int, int]) -> int:
    """The random variable Y: the sum of both dice."""
    return outcome[0] + outcome[1]


# --------------------------------------------------------------------------
# Named distributions: parameters used throughout the lesson and the lab
# --------------------------------------------------------------------------

#: A fair coin, as a Bernoulli(p) reference point.
BERNOULLI_P: float = 0.5

#: Binomial(n, p) used for the Poisson-limit exercise's smallest n, and as a
#: worked example in the lesson.
BINOMIAL_N_EXAMPLE: int = 10
BINOMIAL_P_EXAMPLE: float = 0.3

#: The rate parameter shared by the Poisson distribution and the four
#: Binomial approximations to it in exercise 9. n * p is held fixed at this
#: value as n grows, which is exactly the limiting condition n -> infinity,
#: n * p -> lambda that turns a Binomial into a Poisson.
POISSON_LAMBDA: float = 2.0

#: The four values of n swept in the Poisson-as-Binomial-limit exercise,
#: three decades apart at the ends.
POISSON_LIMIT_NS: tuple[int, ...] = (10, 100, 1_000, 10_000)

#: The range of counts compared between the Binomial and the Poisson pmf.
#: Fifteen values comfortably covers the Poisson(2) distribution's mass --
#: P(X > 14) is under 1e-9 -- so nothing meaningful is cut off.
POISSON_COMPARISON_KS: tuple[int, ...] = tuple(range(0, 15))

#: The rate for the from-scratch exponential sampler.
EXPONENTIAL_RATE: float = 2.0

#: Uniform(0, 0.5) -- the density-above-1 example. Its density is 2
#: everywhere on its support, which is the whole point of exercise 10.
UNIFORM_LOW: float = 0.0
UNIFORM_HIGH: float = 0.5

# --------------------------------------------------------------------------
# Sample sizes and seeds
# --------------------------------------------------------------------------

#: Sample size for the expectation/variance-by-simulation comparison
#: (exercise 3).
EV_SIMULATION_TRIALS: int = 200_000

#: Sample size for the inverse-CDF discrete sampler comparison (exercise 7).
DISCRETE_SAMPLER_TRIALS: int = 200_000

#: Sample size for each of the two exponential samples compared in
#: exercise 8 (from-scratch versus NumPy's own, and the max-gap statistic
#: between their two empirical CDFs).
EXPONENTIAL_SAMPLE_SIZE: int = 50_000

#: The seed every reproducible draw in this lab is built from.
SEED: int = 114

# --------------------------------------------------------------------------
# Tolerances, derived rather than guessed
# --------------------------------------------------------------------------


def standard_error_of_mean(variance: float, n: int) -> float:
    """The standard error of a sample mean: sqrt(variance / n)."""
    return math.sqrt(variance / n)


def standard_error_of_proportion(p: float, n: int) -> float:
    """The standard error of a proportion estimated from n trials."""
    return math.sqrt(p * (1.0 - p) / n)


def dkw_two_sample_threshold(n_a: int, n_b: int, alpha: float = 0.01) -> float:
    """A threshold for the maximum gap between two empirical CDFs drawn
    from the SAME underlying distribution, derived from the
    Dvoretzky-Kiefer-Wolfowitz inequality rather than guessed.

    DKW says: for a sample of size n from a distribution with true CDF F,
    P(sup_x |F_n(x) - F(x)| > eps) <= 2 * exp(-2 * n * eps^2). Solving for
    eps at confidence 1 - alpha (one-sided, so the leading 2 becomes 1)
    gives eps = sqrt(ln(1/alpha) / (2n)). Both empirical CDFs in exercise 8
    are estimating the SAME true exponential CDF, so with probability at
    least 1 - 2*alpha neither one strays more than its own eps from the
    truth, and the gap between them is bounded by the sum of the two eps
    values -- a legitimate, derived threshold, not a number chosen to make
    the test pass.
    """
    eps_a = math.sqrt(math.log(1.0 / alpha) / (2.0 * n_a))
    eps_b = math.sqrt(math.log(1.0 / alpha) / (2.0 * n_b))
    return eps_a + eps_b

"""Exercises 1, 2, 3, 4, 5, 6, 9 and 10: random variables as functions on a
sample space, their pmf/cdf, expectation, variance, and two named
distributions checked numerically.

Every function that can return an exact rational answer returns a
`fractions.Fraction`. Only the named-distribution helpers at the bottom --
which exist to compare a Binomial against a Poisson, and a density against
its integral -- return plain floats, because factorials and exponentials are
not rational.
"""

import math
from fractions import Fraction
from typing import Callable, Iterable, TypeVar

T = TypeVar("T")

# ---------------------------------------------------------------------------
# Exercise 1: the two-dice sum as a random variable, and its pmf
# ---------------------------------------------------------------------------


def dice_sum_pmf() -> dict[int, Fraction]:
    """The probability mass function of Y = sum of two fair dice.

    Built by enumerating all 36 equally likely outcomes and counting how
    many land on each sum -- not by looking up a formula.
    """
    from itertools import product

    counts: dict[int, int] = {}
    for a, b in product(range(1, 7), range(1, 7)):
        total = a + b
        counts[total] = counts.get(total, 0) + 1
    return {k: Fraction(v, 36) for k, v in sorted(counts.items())}


# ---------------------------------------------------------------------------
# Exercise 2: the cdf, as the pmf's running total
# ---------------------------------------------------------------------------


def cdf_from_pmf(pmf: dict[int, Fraction]) -> dict[int, Fraction]:
    """The cumulative distribution function: F(k) = P(X <= k), built by
    accumulating the pmf in increasing order of its keys."""
    running = Fraction(0)
    cdf: dict[int, Fraction] = {}
    for key in sorted(pmf):
        running += pmf[key]
        cdf[key] = running
    return cdf


# ---------------------------------------------------------------------------
# Exercises 3, 4, 5, 6: expectation, variance and covariance, computed
# exactly over any equally-weighted finite space
# ---------------------------------------------------------------------------


def expectation_pmf(pmf: dict[int, Fraction]) -> Fraction:
    """E[X] from a pmf: the weighted average of the values."""
    total = Fraction(0)
    for value, prob in pmf.items():
        total += value * prob
    return total


def variance_pmf(pmf: dict[int, Fraction]) -> Fraction:
    """Var[X] from a pmf: E[(X - E[X])^2]."""
    mean = expectation_pmf(pmf)
    total = Fraction(0)
    for value, prob in pmf.items():
        total += prob * (value - mean) ** 2
    return total


def expectation_over(
    outcomes: Iterable[T], weight: Fraction, func: Callable[[T], Fraction]
) -> Fraction:
    """E[func] over an equally-weighted finite space: sum(weight * func(o)).

    This is the general tool exercises 4, 5 and 6 use: pass it a different
    `func` (the identity, X + Y, X squared, ...) over the SAME 36-outcome
    joint space and it computes the expectation of whatever function you
    hand it, exactly.
    """
    total = Fraction(0)
    for outcome in outcomes:
        total += weight * func(outcome)
    return total


def variance_over(
    outcomes: Iterable[T], weight: Fraction, func: Callable[[T], Fraction]
) -> Fraction:
    """Var[func] over an equally-weighted finite space."""
    outcomes = list(outcomes)
    mean = expectation_over(outcomes, weight, func)
    return expectation_over(outcomes, weight, lambda o: (func(o) - mean) ** 2)


def covariance_over(
    outcomes: Iterable[T],
    weight: Fraction,
    f: Callable[[T], Fraction],
    g: Callable[[T], Fraction],
) -> Fraction:
    """Cov[f, g] = E[(f - E[f]) * (g - E[g])], over the same space."""
    outcomes = list(outcomes)
    mean_f = expectation_over(outcomes, weight, f)
    mean_g = expectation_over(outcomes, weight, g)
    return expectation_over(
        outcomes, weight, lambda o: (f(o) - mean_f) * (g(o) - mean_g)
    )


# ---------------------------------------------------------------------------
# Exercise 9: Binomial and Poisson pmfs, and the gap between them
# ---------------------------------------------------------------------------


def binomial_pmf(n: int, p: float, k: int) -> float:
    """P(K = k) for K ~ Binomial(n, p)."""
    if not 0 <= k <= n:
        return 0.0
    return math.comb(n, k) * (p**k) * ((1.0 - p) ** (n - k))


def poisson_pmf(lam: float, k: int) -> float:
    """P(K = k) for K ~ Poisson(lambda)."""
    if k < 0:
        return 0.0
    return math.exp(-lam) * (lam**k) / math.factorial(k)


def max_binomial_poisson_gap(n: int, p: float, lam: float, ks: Iterable[int]) -> float:
    """The largest |Binomial(n, p).pmf(k) - Poisson(lambda).pmf(k)| over the
    given range of k. As n grows with n * p held at lambda, this gap shrinks
    to zero -- that convergence IS the Poisson limit theorem, and this
    function is what measures it."""
    return max(
        abs(binomial_pmf(n, p, k) - poisson_pmf(lam, k)) for k in ks
    )


# ---------------------------------------------------------------------------
# Exercise 10: a density that exceeds 1, and its integral that does not
# ---------------------------------------------------------------------------


def uniform_density(x: float, low: float, high: float) -> float:
    """The pdf of Uniform(low, high) at x: a constant 1 / (high - low) on
    the support, 0 elsewhere. For Uniform(0, 0.5) that constant is 2 -- a
    density greater than 1, which is legal, because a density is not a
    probability."""
    if low <= x <= high:
        return 1.0 / (high - low)
    return 0.0


def numeric_integral(
    f: Callable[[float], float], low: float, high: float, steps: int
) -> float:
    """The trapezoid-rule numeric integral of f over [low, high], used here
    to confirm that a density integrates to exactly 1 even though its value
    everywhere on the support is 2."""
    if steps < 1:
        raise ValueError("steps must be at least 1")
    width = (high - low) / steps
    total = 0.5 * (f(low) + f(high))
    for i in range(1, steps):
        total += f(low + i * width)
    return total * width

"""Exercises 7 and 8: the inverse-CDF sampling method, written from scratch,
for both a discrete pmf and the exponential distribution -- and a max-gap
statistic to compare two empirical distributions by hand, since scipy is
not installed here.

The idea behind every function in this file is the same one: a single
uniform draw U on (0, 1), pushed through the inverse of a target CDF,
lands on a sample from that target distribution. `numpy.random.Generator`
does exactly this internally for the distributions it knows how to sample;
this file demystifies it by building both halves yourself.
"""

import numpy as np


# ---------------------------------------------------------------------------
# Exercise 7: inverse-CDF sampling for an arbitrary discrete distribution
# ---------------------------------------------------------------------------


def sample_discrete_inverse_cdf(
    pmf: dict[int, float], rng: np.random.Generator, size: int
) -> np.ndarray:
    """Sample from an arbitrary discrete distribution using one uniform
    draw per sample.

    Build the cdf as a running total over the sorted values. Draw U from
    Uniform(0, 1). The sample is the smallest value whose cdf is at least
    U -- geometrically, the point where a vertical line at height U first
    crosses the cdf staircase. `numpy.searchsorted` finds that crossing
    point in one vectorised call instead of a Python loop per draw.
    """
    values = sorted(pmf)
    cumulative = np.cumsum([float(pmf[v]) for v in values])
    # Floating-point summation of many fractions can land at 0.999999999999
    # instead of exactly 1.0. Clamping the last entry to 1.0 guarantees
    # every possible U in [0, 1) has somewhere to land.
    cumulative[-1] = 1.0
    draws = rng.random(size)
    indices = np.searchsorted(cumulative, draws, side="right")
    return np.asarray(values, dtype=float)[indices]


# ---------------------------------------------------------------------------
# Exercise 8: the exponential distribution, sampled from scratch
# ---------------------------------------------------------------------------


def sample_exponential_scratch(
    rate: float, rng: np.random.Generator, size: int
) -> np.ndarray:
    """Sample from Exponential(rate) as -ln(U) / rate.

    The exponential's cdf is F(x) = 1 - exp(-rate * x). Setting F(x) = U and
    solving for x gives x = -ln(1 - U) / rate; since U and 1 - U have the
    same distribution on (0, 1), the simpler -ln(U) / rate is used instead
    -- the standard form of this sampler, and the one worth recognising
    when you meet it again inside a library's source.
    """
    draws = rng.random(size)
    return -np.log(draws) / rate


# ---------------------------------------------------------------------------
# A max-gap statistic between two empirical distributions, written by hand
# because scipy.stats.ks_2samp is not available in this environment
# ---------------------------------------------------------------------------


def empirical_cdf_at(sample: np.ndarray, points: np.ndarray) -> np.ndarray:
    """F_n(x) for every x in `points`: the fraction of `sample` at or below
    each point, read off a sorted copy with a binary search."""
    sorted_sample = np.sort(sample)
    return np.searchsorted(sorted_sample, points, side="right") / len(sample)


def max_gap_statistic(sample_a: np.ndarray, sample_b: np.ndarray) -> float:
    """The largest vertical gap between two empirical CDFs, evaluated at
    every point either sample could change value -- the two-sample
    Kolmogorov-Smirnov statistic, without scipy.

    The maximum can only occur at a point where one of the empirical CDFs
    actually jumps, so evaluating both CDFs at every value that appears in
    either sample (the pooled, sorted set of observations) is exact, not an
    approximation on a grid.
    """
    pooled = np.sort(np.concatenate([sample_a, sample_b]))
    cdf_a = empirical_cdf_at(sample_a, pooled)
    cdf_b = empirical_cdf_at(sample_b, pooled)
    return float(np.max(np.abs(cdf_a - cdf_b)))

"""The reference test suite: real values, real exceptions, no mocking."""

import math
from fractions import Fraction

import numpy as np
import pytest

import dataset as D
import distributions as dist
import sampling as samp

# ---------------------------------------------------------------------------
# dice_sum_pmf
# ---------------------------------------------------------------------------


def test_pmf_has_eleven_entries():
    pmf = dist.dice_sum_pmf()
    assert set(pmf) == set(range(2, 13))


def test_pmf_sums_to_one():
    pmf = dist.dice_sum_pmf()
    assert sum(pmf.values()) == 1


def test_pmf_of_seven_is_one_sixth():
    pmf = dist.dice_sum_pmf()
    assert pmf[7] == Fraction(1, 6)


def test_pmf_returns_fractions():
    pmf = dist.dice_sum_pmf()
    assert all(isinstance(p, Fraction) for p in pmf.values())


def test_pmf_is_symmetric_around_seven():
    pmf = dist.dice_sum_pmf()
    for offset in range(1, 6):
        assert pmf[7 - offset] == pmf[7 + offset]


def test_seven_is_six_times_two():
    pmf = dist.dice_sum_pmf()
    assert pmf[7] == 6 * pmf[2]


def test_the_distribution_is_not_uniform():
    pmf = dist.dice_sum_pmf()
    assert len(set(pmf.values())) > 1


# ---------------------------------------------------------------------------
# cdf_from_pmf
# ---------------------------------------------------------------------------


def test_cdf_is_monotone_non_decreasing():
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    values = [cdf[k] for k in sorted(cdf)]
    assert all(a <= b for a, b in zip(values, values[1:]))


def test_cdf_ends_at_exactly_one():
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    assert cdf[max(cdf)] == 1


def test_cdf_starts_at_the_first_pmf_value():
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    assert cdf[2] == pmf[2]


def test_cdf_difference_equals_pmf_value():
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    assert cdf[7] - cdf[6] == pmf[7]


@pytest.mark.parametrize("k", range(3, 13))
def test_cdf_difference_equals_pmf_at_every_k(k):
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    assert cdf[k] - cdf[k - 1] == pmf[k]


def test_cdf_of_an_interval_matches_a_direct_sum():
    pmf = dist.dice_sum_pmf()
    cdf = dist.cdf_from_pmf(pmf)
    direct = sum((pmf[k] for k in range(5, 10)), Fraction(0))
    assert cdf[9] - cdf[4] == direct


# ---------------------------------------------------------------------------
# expectation_pmf / variance_pmf
# ---------------------------------------------------------------------------


def test_expectation_of_dice_sum_is_seven():
    pmf = dist.dice_sum_pmf()
    assert dist.expectation_pmf(pmf) == 7


def test_variance_of_dice_sum_is_35_over_6():
    pmf = dist.dice_sum_pmf()
    assert dist.variance_pmf(pmf) == Fraction(35, 6)


def test_expectation_of_a_single_die_is_three_point_five():
    pmf = {k: Fraction(1, 6) for k in range(1, 7)}
    assert dist.expectation_pmf(pmf) == Fraction(7, 2)


def test_expectation_need_not_be_an_attainable_value():
    pmf = {k: Fraction(1, 6) for k in range(1, 7)}
    mean = dist.expectation_pmf(pmf)
    assert mean not in pmf


# ---------------------------------------------------------------------------
# expectation_over / variance_over / covariance_over: linearity and its limit
# ---------------------------------------------------------------------------


def test_linearity_holds_for_the_dependent_pair():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    e_x = dist.expectation_over(outcomes, weight, D.first_die)
    e_y = dist.expectation_over(outcomes, weight, D.dice_sum)
    e_sum = dist.expectation_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))
    assert e_sum == e_x + e_y


def test_e_x_is_three_point_five():
    e_x = dist.expectation_over(D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT, D.first_die)
    assert e_x == Fraction(7, 2)


def test_e_y_is_seven():
    e_y = dist.expectation_over(D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT, D.dice_sum)
    assert e_y == 7


def test_variance_of_sum_is_not_the_naive_sum():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    var_x = dist.variance_over(outcomes, weight, D.first_die)
    var_y = dist.variance_over(outcomes, weight, D.dice_sum)
    var_sum = dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))
    assert var_sum != var_x + var_y


def test_variance_of_sum_equals_the_full_covariance_formula():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    var_x = dist.variance_over(outcomes, weight, D.first_die)
    var_y = dist.variance_over(outcomes, weight, D.dice_sum)
    cov_xy = dist.covariance_over(outcomes, weight, D.first_die, D.dice_sum)
    var_sum = dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))
    assert var_sum == var_x + var_y + 2 * cov_xy


def test_covariance_of_x_and_y_is_nonzero():
    cov_xy = dist.covariance_over(D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT, D.first_die, D.dice_sum)
    assert cov_xy != 0
    assert cov_xy == Fraction(35, 12)


def test_covariance_of_independent_like_pair_can_be_zero():
    # First die and "is the second die a 6" behave independently in
    # covariance, since the second die's value is unrelated to the first.
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    cov = dist.covariance_over(outcomes, weight, lambda o: o[0], lambda o: o[1])
    assert cov == 0


# ---------------------------------------------------------------------------
# Jensen's inequality
# ---------------------------------------------------------------------------


def test_jensen_strict_for_a_die():
    outcomes, weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
    e_x2 = dist.expectation_over(outcomes, weight, lambda x: x * x)
    e_x = dist.expectation_over(outcomes, weight, lambda x: x)
    assert e_x2 > e_x**2


def test_jensen_gap_equals_variance():
    outcomes, weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
    e_x2 = dist.expectation_over(outcomes, weight, lambda x: x * x)
    e_x = dist.expectation_over(outcomes, weight, lambda x: x)
    var_x = dist.variance_over(outcomes, weight, lambda x: x)
    assert e_x2 - e_x**2 == var_x


def test_jensen_is_equality_for_a_constant():
    outcomes, weight = (5, 5, 5), Fraction(1, 3)
    e_x2 = dist.expectation_over(outcomes, weight, lambda x: x * x)
    e_x = dist.expectation_over(outcomes, weight, lambda x: x)
    assert e_x2 == e_x**2


# ---------------------------------------------------------------------------
# Var[aX + b]
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("a,b", [(2, 0), (2, 100), (-3, 7), (0.5, -4)])
def test_variance_of_affine_transform(a, b):
    outcomes, weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
    var_x = dist.variance_over(outcomes, weight, lambda x: x)
    var_ax_b = dist.variance_over(outcomes, weight, lambda x: a * x + b)
    assert var_ax_b == a**2 * var_x


# ---------------------------------------------------------------------------
# binomial_pmf / poisson_pmf
# ---------------------------------------------------------------------------


def test_binomial_pmf_sums_to_one():
    n, p = 10, 0.3
    total = sum(dist.binomial_pmf(n, p, k) for k in range(n + 1))
    assert abs(total - 1.0) < 1e-9


def test_binomial_pmf_out_of_range_is_zero():
    assert dist.binomial_pmf(10, 0.3, -1) == 0.0
    assert dist.binomial_pmf(10, 0.3, 11) == 0.0


def test_binomial_pmf_at_zero_successes():
    n, p = 10, 0.3
    assert abs(dist.binomial_pmf(n, p, 0) - (1 - p) ** n) < 1e-12


def test_poisson_pmf_sums_close_to_one_over_a_wide_range():
    lam = 2.0
    total = sum(dist.poisson_pmf(lam, k) for k in range(0, 40))
    assert abs(total - 1.0) < 1e-9


def test_poisson_pmf_negative_k_is_zero():
    assert dist.poisson_pmf(2.0, -1) == 0.0


def test_poisson_pmf_at_zero():
    lam = 2.0
    assert abs(dist.poisson_pmf(lam, 0) - math.exp(-lam)) < 1e-12


# ---------------------------------------------------------------------------
# max_binomial_poisson_gap: the Poisson-as-Binomial-limit convergence
# ---------------------------------------------------------------------------


def test_poisson_limit_gap_shrinks_monotonically():
    lam = D.POISSON_LAMBDA
    gaps = [
        dist.max_binomial_poisson_gap(n, lam / n, lam, D.POISSON_COMPARISON_KS)
        for n in D.POISSON_LIMIT_NS
    ]
    assert all(a > b for a, b in zip(gaps, gaps[1:]))


def test_poisson_limit_gap_at_largest_n_is_tiny():
    lam = D.POISSON_LAMBDA
    n = D.POISSON_LIMIT_NS[-1]
    gap = dist.max_binomial_poisson_gap(n, lam / n, lam, D.POISSON_COMPARISON_KS)
    assert gap < 1e-3


# ---------------------------------------------------------------------------
# uniform_density / numeric_integral
# ---------------------------------------------------------------------------


def test_uniform_density_is_two_on_the_support():
    assert dist.uniform_density(0.25, 0.0, 0.5) == 2.0


def test_uniform_density_exceeds_one():
    assert dist.uniform_density(0.1, 0.0, 0.5) > 1.0


def test_uniform_density_is_zero_outside_support():
    assert dist.uniform_density(0.6, 0.0, 0.5) == 0.0
    assert dist.uniform_density(-0.1, 0.0, 0.5) == 0.0


def test_numeric_integral_of_the_density_is_one():
    integral = dist.numeric_integral(lambda x: dist.uniform_density(x, 0.0, 0.5), 0.0, 0.5, 50_000)
    assert abs(integral - 1.0) < 1e-6


def test_numeric_integral_of_a_constant_one():
    integral = dist.numeric_integral(lambda x: 1.0, 0.0, 1.0, 1000)
    assert abs(integral - 1.0) < 1e-9


def test_numeric_integral_rejects_zero_steps():
    with pytest.raises(ValueError):
        dist.numeric_integral(lambda x: 1.0, 0.0, 1.0, 0)


# ---------------------------------------------------------------------------
# sample_discrete_inverse_cdf
# ---------------------------------------------------------------------------


def test_discrete_sampler_returns_only_values_in_the_pmf():
    pmf = {k: float(v) for k, v in dist.dice_sum_pmf().items()}
    rng = np.random.default_rng(0)
    draws = samp.sample_discrete_inverse_cdf(pmf, rng, 2_000)
    assert set(draws.astype(int).tolist()) <= set(pmf)


def test_discrete_sampler_empirical_frequencies_are_close():
    pmf = {k: float(v) for k, v in dist.dice_sum_pmf().items()}
    rng = np.random.default_rng(1)
    n = 100_000
    draws = samp.sample_discrete_inverse_cdf(pmf, rng, n)
    values, counts = np.unique(draws, return_counts=True)
    empirical = dict(zip(values.astype(int), counts / n))
    for k, p in pmf.items():
        se = D.standard_error_of_proportion(p, n)
        assert abs(empirical.get(k, 0.0) - p) < 4.0 * se


def test_discrete_sampler_same_seed_is_reproducible():
    pmf = {k: float(v) for k, v in dist.dice_sum_pmf().items()}
    a = samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(7), 500)
    b = samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(7), 500)
    assert np.array_equal(a, b)


def test_discrete_sampler_different_seed_differs():
    pmf = {k: float(v) for k, v in dist.dice_sum_pmf().items()}
    a = samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(7), 500)
    b = samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(8), 500)
    assert not np.array_equal(a, b)


def test_discrete_sampler_handles_a_skewed_two_point_pmf():
    pmf = {0: 0.9, 1: 0.1}
    rng = np.random.default_rng(2)
    n = 50_000
    draws = samp.sample_discrete_inverse_cdf(pmf, rng, n)
    empirical_one = (draws == 1).mean()
    se = D.standard_error_of_proportion(0.1, n)
    assert abs(empirical_one - 0.1) < 4.0 * se


# ---------------------------------------------------------------------------
# sample_exponential_scratch
# ---------------------------------------------------------------------------


def test_exponential_scratch_mean_near_one_over_rate():
    rate = 2.0
    rng = np.random.default_rng(3)
    n = 100_000
    draws = samp.sample_exponential_scratch(rate, rng, n)
    tol = 3.0 * D.standard_error_of_mean((1.0 / rate) ** 2, n)
    assert abs(draws.mean() - 1.0 / rate) < tol


def test_exponential_scratch_is_never_negative():
    rng = np.random.default_rng(4)
    draws = samp.sample_exponential_scratch(1.0, rng, 10_000)
    assert (draws >= 0).all()


def test_exponential_scratch_matches_builtin_mean_within_tolerance():
    rate = 2.0
    n = 50_000
    rng = np.random.default_rng(5)
    scratch = samp.sample_exponential_scratch(rate, rng, n)
    built_in = rng.exponential(scale=1.0 / rate, size=n)
    tol = 3.0 * D.standard_error_of_mean((1.0 / rate) ** 2, n)
    assert abs(scratch.mean() - built_in.mean()) < 2 * tol


# ---------------------------------------------------------------------------
# empirical_cdf_at / max_gap_statistic
# ---------------------------------------------------------------------------


def test_empirical_cdf_is_between_zero_and_one():
    sample = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    points = np.array([0.0, 2.5, 10.0])
    cdf_vals = samp.empirical_cdf_at(sample, points)
    assert (cdf_vals >= 0).all() and (cdf_vals <= 1).all()
    assert cdf_vals[0] == 0.0
    assert cdf_vals[-1] == 1.0


def test_max_gap_is_zero_for_identical_samples():
    sample = np.array([1.0, 2.0, 3.0, 4.0])
    assert samp.max_gap_statistic(sample, sample) == 0.0


def test_max_gap_is_one_for_disjoint_supports():
    a = np.array([0.0, 0.0, 0.0])
    b = np.array([100.0, 100.0, 100.0])
    assert samp.max_gap_statistic(a, b) == 1.0


def test_max_gap_between_two_exponential_samples_is_small():
    rng = np.random.default_rng(9)
    n = 30_000
    a = samp.sample_exponential_scratch(2.0, rng, n)
    b = rng.exponential(scale=0.5, size=n)
    gap = samp.max_gap_statistic(a, b)
    threshold = D.dkw_two_sample_threshold(n, n)
    assert gap < threshold


# ---------------------------------------------------------------------------
# dkw_two_sample_threshold / standard_error helpers
# ---------------------------------------------------------------------------


def test_dkw_threshold_shrinks_as_n_grows():
    small = D.dkw_two_sample_threshold(1_000, 1_000)
    large = D.dkw_two_sample_threshold(100_000, 100_000)
    assert large < small


def test_standard_error_of_mean_scales_as_inverse_sqrt_n():
    se_100 = D.standard_error_of_mean(4.0, 100)
    se_10000 = D.standard_error_of_mean(4.0, 10_000)
    assert abs(se_100 / se_10000 - 10.0) < 1e-9


def test_standard_error_of_proportion_is_symmetric_in_p():
    assert D.standard_error_of_proportion(0.3, 1000) == D.standard_error_of_proportion(0.7, 1000)

lesson_id: D114
day: 114
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-114-random-variables-and-distributions
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
  - .venv/bin/python3 -c "import numpy; print(numpy.__version__)"
run_commands:
  - 'cd examples && ../.venv/bin/python3 01_pmf_of_a_sum.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_cdf_from_pmf.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_expectation_and_variance.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_linearity_with_dependence.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_variance_is_not_additive.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 06_jensens_inequality.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 07_inverse_cdf_discrete_sampler.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 08_exponential_from_scratch.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 09_poisson_as_binomial_limit.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 10_density_above_one.py && cd ..'
  - .venv/bin/pytest examples -q -p no:cacheprovider
  - .venv/bin/pytest starter -q -p no:cacheprovider
test_commands:
  - bash tests/run_tests.sh
cleanup_commands:
  - "find . -path ./.venv -prune -o -type d -name '__pycache__' -print -exec rm -rf -- {} +"
  - rm -rf .pytest_cache
  - 'rm -rf .venv  # optional: removes the lab virtual environment'
  - 'git checkout -- starter/  # optional: reset your work'
requires_network: true
requires_api_key: false
estimated_minutes: 35
last_executed: '2026-08-17'
executed_on: 'macOS 26.5.2 (Apple Silicon, arm64), Python 3.14.0, numpy 2.5.2, pytest 9.1.1, bash 3.2.57 -- bash tests/run_tests.sh -> 63 checks, 0 failure(s), exit 0; pytest examples -> 69 passed; pytest starter -> 2 passed, 43 skipped on an untouched checkout, and 45 passed against a fully solved copy of starter/ (verified by temporarily copying the reference distributions.py and sampling.py into starter/ and filling in every answers.py prediction, then restoring the blank skeletons -- the skip counts before and after that restore were confirmed identical, 43 both times, and collecting both suites together also reports 43 skipped, proving the two conftest.py import guards work). All ten reference scripts exit 0 with every internal assertion holding. Everything was run through a real lab-local .venv created by the documented setup commands, not through an authoring environment; pip install used the network exactly once, as documented. Section 6 of the harness re-runs the reference pytest suite with the variance-non-additivity assertion temporarily replaced by the WRONG belief that Var[X+Y] equals Var[X]+Var[Y], confirms the run exits non-zero with exactly one named failure, and restores the original file -- so the suite is demonstrated to be capable of failing rather than merely claimed to be. Separately from the harness''s own self-test, one of the ten reference scripts (01_pmf_of_a_sum.py) was manually edited to assert pmf[7] == Fraction(1, 7) instead of Fraction(1, 6), the full harness was re-run and confirmed to exit 1 with 2 named failures, and the file was restored and the harness re-confirmed green at 63 checks, 0 failures. Three honesty notes from this run. FIRST: scipy, pandas and matplotlib are not installed in this environment. scipy.stats is described from its public documentation in the lesson''s Tools section and explicitly marked as not run here; no output attributed to it anywhere in this lab or its lesson was actually produced by it. SECOND: every sampled figure in this lab (the simulated dice-sum mean, the inverse-CDF sampler''s empirical frequencies, both exponential sample means, and the max-gap statistic between their empirical cdfs) is a freshly measured number rather than a fixed literal, checked against a tolerance derived from a standard error or the Dvoretzky-Kiefer-Wolfowitz inequality rather than a value chosen to make the test pass; the specific figures will differ slightly on another machine or NumPy version, as documented in expected-output/FIELDS.md. THIRD: the max-gap statistic between the from-scratch exponential sampler and NumPy''s own (0.0050 on this run) was checked across five different seeds during development (0.0043 to 0.0058) and stayed comfortably under the DKW-derived threshold (0.01357) every time, but this is reported as observed behaviour across a handful of runs, not as a formal proof the threshold can never be exceeded. Every exact Fraction claim in this lab (the dice-sum pmf, the cdf, the expectation and variance identities, the linearity and non-additivity results, and Jensen''s inequality) is exact rational arithmetic and is therefore identical on any correct Python implementation, anywhere.'

# What is installed, why, and what it costs

Two packages, both free and open source, both installed into a lab-local
virtual environment that `rm -rf .venv` completely undoes.

| Package | Version pinned | Licence | What this lab uses it for |
| --- | --- | --- | --- |
| `numpy` | 2.5.2 | BSD 3-Clause | The `numpy.random.Generator` built by `default_rng(seed)`, used for every simulation and every sampler in exercises 3, 7, 8 and 9. |
| `pytest` | 9.1.1 | MIT | The reference suite (69 tests) and your running score in `starter/`. |

There is no paid tier of anything in this lab, no account, no key and no
signup, personally or commercially.

## The one time the network is needed

```bash
.venv/bin/pip install -r requirements/requirements.txt
```

That is the only command in the lab that opens a connection. Section 7 of
`tests/run_tests.sh` greps every source file in `examples/` and `starter/`
to prove that nothing else does.

## If you cannot install anything at all

Exercises 1, 2, 4, 5, 6, 9 and 10 need only `fractions` and `math` from the
standard library and do not touch NumPy at all. Only exercises 3, 7 and 8
need `numpy.random.Generator`. If NumPy is unavailable, Python's own
`random` module can stand in for a uniform draw:

```python
import random

def sample_exponential_scratch_stdlib(rate: float, rng: random.Random, size: int) -> list[float]:
    import math
    return [-math.log(rng.random()) / rate for _ in range(size)]
```

It is slower -- a Python-level loop instead of a vectorised NumPy call --
but it is the same statistics. What you lose is `pytest`, so no running
score and no skip-versus-fail distinction; you would read the numbers back
yourself instead.

## What is deliberately *not* installed

`scipy.stats` does this job too -- its `rv_continuous` and `rv_discrete`
base classes are the shape every named distribution in this lesson's table
would map onto, with `.pmf()`/`.pdf()`, `.cdf()`, `.mean()`, `.var()` and
`.rvs()` methods -- and it is **not installed in this environment, and no
output from it is reproduced anywhere** in this lab or its lesson. The
lesson's Tools section describes it from its documentation and marks it as
not run here.

`pandas` and `matplotlib` are also not installed and are not needed for
anything in this lab; every table here is small enough to enumerate and
print directly, and every plot the lesson describes is described in prose
rather than rendered.

That is not a limitation to apologise for. Every exact claim this lab makes
is computed with `fractions.Fraction` over a finite enumerated space, and
every sampled claim is drawn from the standard library and NumPy's random
module and checked against a derived tolerance -- nothing here depends on
a statistics package to be correct.

numpy==2.5.2
pytest==9.1.1

# The ten exercises

Work through these in order. Predict the answer to each `answers.py`
question *before* running anything — the ratio in exercise 1 and the
variance-non-additivity in exercise 5 only catch you if you commit to a
guess first.

Check yourself as you go:

```bash
.venv/bin/pytest starter -q
```

Unattempted work reports as **skipped**, never failed. Wrong work **fails**
with your answer printed beside the correct one.

## 1. The pmf of a sum (`distributions.py`)

`dice_sum_pmf()`, built by enumerating all 36 outcomes of two dice with
`itertools.product` and counting how many land on each sum. Return exact
`fractions.Fraction` values. Assert `pmf[7] == Fraction(1, 6)` and that 7 is
exactly six times as likely as 2 or 12.

## 2. The cdf as a running total (`distributions.py`)

`cdf_from_pmf(pmf)`. Assert it is monotone non-decreasing, ends at exactly
1, and that `cdf[7] - cdf[6] == pmf[7]` exactly.

## 3. Expectation and variance (`distributions.py`)

`expectation_pmf(pmf)` and `variance_pmf(pmf)`, computed from the
definition. Compared in the reference script against a large seeded
simulation using `numpy.random.default_rng` and the standard library's
`statistics` module.

## 4 and 5. Linearity and non-additivity (`distributions.py`)

`expectation_over`, `variance_over` and `covariance_over` — three general
tools that work over any equally-weighted finite space. Let X be the first
die and Y the sum of both dice; Y depends on X directly. Assert
`E[X+Y] == E[X] + E[Y]` exactly (exercise 4), then assert `Var[X+Y] !=
Var[X] + Var[Y]` but `Var[X+Y] == Var[X] + Var[Y] + 2*Cov(X,Y)` exactly
(exercise 5).

## 6. Jensen's inequality (`distributions.py`)

Using the same three tools on a single die: assert `E[X^2] > (E[X])^2`
exactly, and that the gap equals `Var[X]` exactly.

## 7. An inverse-CDF sampler for a discrete pmf (`sampling.py`)

`sample_discrete_inverse_cdf(pmf, rng, size)`, written from scratch with
one uniform draw per sample and `numpy.searchsorted` against the pmf's own
cdf. Assert the empirical frequencies match the pmf within a tolerance
derived from the standard error of a proportion, and that the same seed
reproduces identical draws.

## 8. The exponential distribution from scratch (`sampling.py`)

`sample_exponential_scratch(rate, rng, size)` as `-ln(U) / rate`. Compared
against `Generator.exponential` on sample mean, and with a max-gap
statistic between the two empirical cdfs that you also write by hand
(`empirical_cdf_at`, `max_gap_statistic`) — scipy is not installed, so no
`ks_2samp`.

## 9. Poisson as a Binomial limit (`distributions.py`)

`binomial_pmf(n, p, k)`, `poisson_pmf(lam, k)` and
`max_binomial_poisson_gap(n, p, lam, ks)`. With lambda held at 2 and p =
lambda / n, assert the maximum pmf gap decreases monotonically as n grows
across 10, 100, 1,000 and 10,000.

## 10. A density that exceeds 1 (`distributions.py`)

`uniform_density(x, low, high)` and `numeric_integral(f, low, high,
steps)`. For Uniform(0, 0.5), assert the density equals exactly 2 while its
numeric integral over the support equals 1 to a small tolerance.

"""Exercises 1 through 10 -- eighteen predictions.

Replace each `None` with the value you think is correct. A `None` is a
skip, not a failure: `pytest starter -q` counts only what you have
attempted. When you are wrong it prints both your answer and the real one,
so a wrong guess is worth more than a blank.

Predict BEFORE you run anything. The two that catch almost everyone are the
seven-to-two ratio (exercise 1) and whether Var[X+Y] equals the naive sum
(exercise 5) -- and they only catch you if you commit first.

Every answer is a number or a Python bool.
"""

ANSWERS: dict[str, object] = {
    # ----------------------------------------------------------------------
    # Exercise 1 -- the pmf of a sum
    # ----------------------------------------------------------------------
    # 1.1 P(the two dice sum to 7), as a decimal.
    "p_sum_seven": None,
    # 1.2 How many times as likely is a sum of 7 as a sum of 2?
    "ratio_seven_to_two": None,
    # ----------------------------------------------------------------------
    # Exercise 2 -- the cdf
    # ----------------------------------------------------------------------
    # 2.1 F(12), the cdf at its largest value.
    "cdf_at_twelve": None,
    # 2.2 F(7) - F(6), as a decimal.
    "cdf_difference_seven_six": None,
    # ----------------------------------------------------------------------
    # Exercise 3 -- expectation and variance
    # ----------------------------------------------------------------------
    # 3.1 E[Y], the expected sum of two dice.
    "expectation_of_sum": None,
    # 3.2 Var[Y], as a decimal.
    "variance_of_sum": None,
    # ----------------------------------------------------------------------
    # Exercise 4 -- linearity with a dependent pair
    # ----------------------------------------------------------------------
    # 4.1 E[X + Y] where X = first die, Y = sum of both dice.
    "expectation_x_plus_y": None,
    # 4.2 Does E[X + Y] == E[X] + E[Y] hold even though X and Y are
    #     dependent?
    "linearity_holds_for_dependent_pair": None,
    # ----------------------------------------------------------------------
    # Exercise 5 -- variance is not additive
    # ----------------------------------------------------------------------
    # 5.1 Does Var[X + Y] == Var[X] + Var[Y] (the naive, WRONG sum)?
    "variance_naive_sum_holds": None,
    # 5.2 Var[X + Y], as a decimal.
    "variance_x_plus_y": None,
    # ----------------------------------------------------------------------
    # Exercise 6 -- Jensen's inequality
    # ----------------------------------------------------------------------
    # 6.1 Is E[X^2] strictly greater than (E[X])^2 for a single die?
    "jensen_strict_for_die": None,
    # 6.2 Does the gap E[X^2] - (E[X])^2 equal Var[X] exactly?
    "jensen_gap_equals_variance": None,
    # ----------------------------------------------------------------------
    # Exercise 7 -- inverse-CDF discrete sampling
    # ----------------------------------------------------------------------
    # 7.1 Does the same seed reproduce identical draws?
    "discrete_sampler_reproducible": None,
    # ----------------------------------------------------------------------
    # Exercise 8 -- exponential from scratch
    # ----------------------------------------------------------------------
    # 8.1 Are all draws from the from-scratch exponential sampler
    #     non-negative?
    "exponential_scratch_nonnegative": None,
    # ----------------------------------------------------------------------
    # Exercise 9 -- Poisson as a Binomial limit
    # ----------------------------------------------------------------------
    # 9.1 Does the maximum pmf gap decrease monotonically as n grows across
    #     10, 100, 1000, 10000?
    "poisson_gap_decreases_monotonically": None,
    # ----------------------------------------------------------------------
    # Exercise 10 -- density above 1
    # ----------------------------------------------------------------------
    # 10.1 The density of Uniform(0, 0.5) at any point in its support.
    "uniform_density_value": None,
    # 10.2 Is that density value greater than 1?
    "uniform_density_exceeds_one": None,
    # 10.3 Does the numeric integral of that density over its support equal
    #      1 (to a small tolerance)?
    "uniform_integral_equals_one": None,
}

"""Make this directory's own modules the ones its tests import.

Both `examples/` and `starter/` contain modules called `distributions`,
`sampling` and `dataset`, and pytest imports test files by putting their
directory on `sys.path`. Without this file, running `pytest` across both
directories at once would import whichever `distributions` was seen first
and reuse it for the other suite -- so the starter tests would silently
pass against the reference solution instead of skipping. That is a wrong
answer with a green tick on it, which is the worst kind.

So: put this directory first on the import path, and drop any
already-imported module of those names that came from somewhere else.
"""

import sys
from pathlib import Path

HERE = str(Path(__file__).parent.resolve())

if HERE in sys.path:
    sys.path.remove(HERE)
sys.path.insert(0, HERE)

for name in ("distributions", "sampling", "dataset", "answers"):
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", "") or ""
    if module is not None and not origin.startswith(HERE):
        del sys.modules[name]

"""The sample spaces, the named-distribution parameters, and every tolerance
this lab compares against.

Read this file. Nothing here is tuned: every exact figure below is either
computed by enumeration (`itertools.product` and `fractions.Fraction`) or
derived algebraically and then checked against the enumeration in the tests.
Every simulation tolerance is derived from a standard error written out
beside it -- never chosen by running a test and loosening a number until it
passed.
"""

import itertools
import math
from fractions import Fraction

# --------------------------------------------------------------------------
# The two-dice sample space, carried over from Day 113 and reused here as
# the running example for a random variable -- the function "sum of the two
# faces" mapping this 36-outcome sample space to the integers.
# --------------------------------------------------------------------------

DIE_FACES: tuple[int, ...] = tuple(range(1, 7))

#: Every ordered pair (first die, second die), 36 equally likely outcomes.
TWO_DICE_SPACE: tuple[tuple[int, int], ...] = tuple(
    itertools.product(DIE_FACES, DIE_FACES)
)
assert len(TWO_DICE_SPACE) == 36

#: The weight of a single outcome in an equally-likely 36-outcome space.
TWO_DICE_WEIGHT: Fraction = Fraction(1, 36)

#: The weight of a single outcome in an equally-likely 6-outcome space (one
#: die alone), used for the Jensen's-inequality exercise.
ONE_DIE_WEIGHT: Fraction = Fraction(1, 6)


def first_die(outcome: tuple[int, int]) -> int:
    """The random variable X: the first die's face."""
    return outcome[0]


def dice_sum(outcome: tuple[int, int]) -> int:
    """The random variable Y: the sum of both dice."""
    return outcome[0] + outcome[1]


# --------------------------------------------------------------------------
# Named distributions: parameters used throughout the lesson and the lab
# --------------------------------------------------------------------------

#: A fair coin, as a Bernoulli(p) reference point.
BERNOULLI_P: float = 0.5

#: Binomial(n, p) used for the Poisson-limit exercise's smallest n, and as a
#: worked example in the lesson.
BINOMIAL_N_EXAMPLE: int = 10
BINOMIAL_P_EXAMPLE: float = 0.3

#: The rate parameter shared by the Poisson distribution and the four
#: Binomial approximations to it in exercise 9. n * p is held fixed at this
#: value as n grows, which is exactly the limiting condition n -> infinity,
#: n * p -> lambda that turns a Binomial into a Poisson.
POISSON_LAMBDA: float = 2.0

#: The four values of n swept in the Poisson-as-Binomial-limit exercise,
#: three decades apart at the ends.
POISSON_LIMIT_NS: tuple[int, ...] = (10, 100, 1_000, 10_000)

#: The range of counts compared between the Binomial and the Poisson pmf.
#: Fifteen values comfortably covers the Poisson(2) distribution's mass --
#: P(X > 14) is under 1e-9 -- so nothing meaningful is cut off.
POISSON_COMPARISON_KS: tuple[int, ...] = tuple(range(0, 15))

#: The rate for the from-scratch exponential sampler.
EXPONENTIAL_RATE: float = 2.0

#: Uniform(0, 0.5) -- the density-above-1 example. Its density is 2
#: everywhere on its support, which is the whole point of exercise 10.
UNIFORM_LOW: float = 0.0
UNIFORM_HIGH: float = 0.5

# --------------------------------------------------------------------------
# Sample sizes and seeds
# --------------------------------------------------------------------------

#: Sample size for the expectation/variance-by-simulation comparison
#: (exercise 3).
EV_SIMULATION_TRIALS: int = 200_000

#: Sample size for the inverse-CDF discrete sampler comparison (exercise 7).
DISCRETE_SAMPLER_TRIALS: int = 200_000

#: Sample size for each of the two exponential samples compared in
#: exercise 8 (from-scratch versus NumPy's own, and the max-gap statistic
#: between their two empirical CDFs).
EXPONENTIAL_SAMPLE_SIZE: int = 50_000

#: The seed every reproducible draw in this lab is built from.
SEED: int = 114

# --------------------------------------------------------------------------
# Tolerances, derived rather than guessed
# --------------------------------------------------------------------------


def standard_error_of_mean(variance: float, n: int) -> float:
    """The standard error of a sample mean: sqrt(variance / n)."""
    return math.sqrt(variance / n)


def standard_error_of_proportion(p: float, n: int) -> float:
    """The standard error of a proportion estimated from n trials."""
    return math.sqrt(p * (1.0 - p) / n)


def dkw_two_sample_threshold(n_a: int, n_b: int, alpha: float = 0.01) -> float:
    """A threshold for the maximum gap between two empirical CDFs drawn
    from the SAME underlying distribution, derived from the
    Dvoretzky-Kiefer-Wolfowitz inequality rather than guessed.

    DKW says: for a sample of size n from a distribution with true CDF F,
    P(sup_x |F_n(x) - F(x)| > eps) <= 2 * exp(-2 * n * eps^2). Solving for
    eps at confidence 1 - alpha (one-sided, so the leading 2 becomes 1)
    gives eps = sqrt(ln(1/alpha) / (2n)). Both empirical CDFs in exercise 8
    are estimating the SAME true exponential CDF, so with probability at
    least 1 - 2*alpha neither one strays more than its own eps from the
    truth, and the gap between them is bounded by the sum of the two eps
    values -- a legitimate, derived threshold, not a number chosen to make
    the test pass.
    """
    eps_a = math.sqrt(math.log(1.0 / alpha) / (2.0 * n_a))
    eps_b = math.sqrt(math.log(1.0 / alpha) / (2.0 * n_b))
    return eps_a + eps_b

"""Exercises 1, 2, 3, 4, 5, 6, 9 and 10: random variables as functions on a
sample space, their pmf/cdf, expectation, variance, and two named
distributions checked numerically.

Fill in every function below. Read `dataset.py` first -- nothing in it
needs to change. Every function that can return an exact rational answer
should return a `fractions.Fraction`.
"""

import math
from fractions import Fraction
from typing import Callable, Iterable, TypeVar

T = TypeVar("T")

# ---------------------------------------------------------------------------
# Exercise 1: the two-dice sum as a random variable, and its pmf
# ---------------------------------------------------------------------------


def dice_sum_pmf() -> dict[int, Fraction]:
    """The probability mass function of Y = sum of two fair dice.

    Enumerate all 36 equally likely (first die, second die) outcomes with
    `itertools.product(range(1, 7), range(1, 7))`, count how many land on
    each sum, and return {sum: Fraction(count, 36)}.
    """
    raise NotImplementedError


# ---------------------------------------------------------------------------
# Exercise 2: the cdf, as the pmf's running total
# ---------------------------------------------------------------------------


def cdf_from_pmf(pmf: dict[int, Fraction]) -> dict[int, Fraction]:
    """F(k) = P(X <= k), built by accumulating the pmf in increasing order
    of its keys."""
    raise NotImplementedError


# ---------------------------------------------------------------------------
# Exercises 3, 4, 5, 6: expectation, variance and covariance, computed
# exactly over any equally-weighted finite space
# ---------------------------------------------------------------------------


def expectation_pmf(pmf: dict[int, Fraction]) -> Fraction:
    """E[X] from a pmf: the weighted average of the values."""
    raise NotImplementedError


def variance_pmf(pmf: dict[int, Fraction]) -> Fraction:
    """Var[X] from a pmf: E[(X - E[X])^2]."""
    raise NotImplementedError


def expectation_over(
    outcomes: Iterable[T], weight: Fraction, func: Callable[[T], Fraction]
) -> Fraction:
    """E[func] over an equally-weighted finite space: sum(weight * func(o))
    over every outcome o. `weight` is the same Fraction for every outcome,
    since the space is equally likely."""
    raise NotImplementedError


def variance_over(
    outcomes: Iterable[T], weight: Fraction, func: Callable[[T], Fraction]
) -> Fraction:
    """Var[func] over an equally-weighted finite space. Compute the mean
    with `expectation_over` first, then the expectation of the squared
    deviation from it."""
    raise NotImplementedError


def covariance_over(
    outcomes: Iterable[T],
    weight: Fraction,
    f: Callable[[T], Fraction],
    g: Callable[[T], Fraction],
) -> Fraction:
    """Cov[f, g] = E[(f - E[f]) * (g - E[g])], over the same space."""
    raise NotImplementedError


# ---------------------------------------------------------------------------
# Exercise 9: Binomial and Poisson pmfs, and the gap between them
# ---------------------------------------------------------------------------


def binomial_pmf(n: int, p: float, k: int) -> float:
    """P(K = k) for K ~ Binomial(n, p). Return 0.0 if k is outside [0, n].
    `math.comb(n, k)` gives the binomial coefficient."""
    raise NotImplementedError


def poisson_pmf(lam: float, k: int) -> float:
    """P(K = k) for K ~ Poisson(lambda). Return 0.0 if k < 0."""
    raise NotImplementedError


def max_binomial_poisson_gap(n: int, p: float, lam: float, ks: Iterable[int]) -> float:
    """The largest |Binomial(n, p).pmf(k) - Poisson(lambda).pmf(k)| over the
    given range of k."""
    raise NotImplementedError


# ---------------------------------------------------------------------------
# Exercise 10: a density that exceeds 1, and its integral that does not
# ---------------------------------------------------------------------------


def uniform_density(x: float, low: float, high: float) -> float:
    """The pdf of Uniform(low, high) at x: 1 / (high - low) on the
    support, 0 elsewhere."""
    raise NotImplementedError


def numeric_integral(
    f: Callable[[float], float], low: float, high: float, steps: int
) -> float:
    """The trapezoid-rule numeric integral of f over [low, high]. Raise
    ValueError if steps < 1. The trapezoid rule: split [low, high] into
    `steps` equal-width panels, weight the two endpoint evaluations by 0.5,
    weight every interior evaluation by 1, and multiply the total by the
    panel width."""
    raise NotImplementedError

"""Exercises 7 and 8: the inverse-CDF sampling method, written from scratch,
for both a discrete pmf and the exponential distribution -- and a max-gap
statistic to compare two empirical distributions by hand, since scipy is
not installed here.

Fill in every function below.
"""

import numpy as np


# ---------------------------------------------------------------------------
# Exercise 7: inverse-CDF sampling for an arbitrary discrete distribution
# ---------------------------------------------------------------------------


def sample_discrete_inverse_cdf(
    pmf: dict[int, float], rng: np.random.Generator, size: int
) -> np.ndarray:
    """Sample from an arbitrary discrete distribution using one uniform
    draw per sample.

    Build the cdf as a running total over the sorted values (clamp the last
    entry to exactly 1.0 to guard against floating-point summation landing
    a hair below it). Draw `size` values from `rng.random(size)`. Use
    `numpy.searchsorted(cumulative, draws, side="right")` to find, for each
    draw, the index of the smallest cdf value at least as large as the
    draw. Index into the sorted values with that array of indices.
    """
    raise NotImplementedError


# ---------------------------------------------------------------------------
# Exercise 8: the exponential distribution, sampled from scratch
# ---------------------------------------------------------------------------


def sample_exponential_scratch(
    rate: float, rng: np.random.Generator, size: int
) -> np.ndarray:
    """Sample from Exponential(rate) as -ln(U) / rate, where U is drawn
    from `rng.random(size)`."""
    raise NotImplementedError


# ---------------------------------------------------------------------------
# A max-gap statistic between two empirical distributions, written by hand
# because scipy.stats.ks_2samp is not available in this environment
# ---------------------------------------------------------------------------


def empirical_cdf_at(sample: np.ndarray, points: np.ndarray) -> np.ndarray:
    """F_n(x) for every x in `points`: the fraction of `sample` at or below
    each point. Sort a copy of `sample`, then use
    `numpy.searchsorted(sorted_sample, points, side="right") / len(sample)`.
    """
    raise NotImplementedError


def max_gap_statistic(sample_a: np.ndarray, sample_b: np.ndarray) -> float:
    """The largest vertical gap between two empirical CDFs, evaluated at
    every point either sample could change value: pool and sort both
    samples, evaluate both empirical cdfs at every pooled point with
    `empirical_cdf_at`, and return the maximum absolute difference."""
    raise NotImplementedError

"""Your running score. Unattempted work SKIPS; wrong work FAILS with both values.

Run from the lab directory:

    .venv/bin/pytest starter -q

On an untouched checkout this reports one pass and everything else skipped.
A skip means "not attempted". A failure means "attempted and wrong", and the
message shows your answer next to the real one so you can see the gap
rather than guess at it.

Nothing in here checks that a function exists or that a file is present.
Every test runs your code and compares a value.
"""

from fractions import Fraction

import numpy as np
import pytest

import answers
import dataset as D
import distributions as dist
import sampling as samp

# --------------------------------------------------------------------------
# The skip machinery
# --------------------------------------------------------------------------


def need(value, what):
    """Skip if the exercise has not been attempted, otherwise hand it back."""
    if value is None:
        pytest.skip(f"not attempted yet: {what}")
    return value


def attempt(fn, what):
    """Call something that may not be written yet, and skip if it is not."""
    try:
        result = fn()
    except (TypeError, AttributeError, NotImplementedError):
        pytest.skip(f"not attempted yet: {what}")
    if result is None:
        pytest.skip(f"not attempted yet: {what}")
    return result


def close(got, want, tol, what):
    assert abs(float(got) - float(want)) < tol, (
        f"{what}: your answer {got!r}, expected {want!r} "
        f"(difference {abs(float(got) - float(want)):.3e}, tolerance {tol:g})"
    )


def test_the_suite_itself_runs():
    """One test that always passes, so a green run is distinguishable from
    a collection error that quietly ran nothing at all."""
    assert len(D.TWO_DICE_SPACE) == 36


# --------------------------------------------------------------------------
# Exercise 1 -- the pmf of a sum
# --------------------------------------------------------------------------


def test_1_pmf_has_eleven_entries():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    assert set(pmf) == set(range(2, 13))


def test_1_pmf_of_seven_is_exactly_one_sixth():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    assert pmf[7] == Fraction(1, 6), "pmf[7] must be exactly Fraction(1, 6)"


def test_1_pmf_returns_fractions():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    assert all(isinstance(p, Fraction) for p in pmf.values()), (
        "every pmf value must be a Fraction, not a float"
    )


def test_1_seven_is_six_times_two():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    assert pmf[7] == 6 * pmf[2]


# --------------------------------------------------------------------------
# Exercise 2 -- the cdf
# --------------------------------------------------------------------------


def test_2_cdf_is_monotone():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    cdf = attempt(lambda: dist.cdf_from_pmf(pmf), "cdf_from_pmf")
    values = [cdf[k] for k in sorted(cdf)]
    assert all(a <= b for a, b in zip(values, values[1:]))


def test_2_cdf_ends_at_one():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    cdf = attempt(lambda: dist.cdf_from_pmf(pmf), "cdf_from_pmf")
    assert cdf[max(cdf)] == 1


def test_2_cdf_difference_equals_pmf():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    cdf = attempt(lambda: dist.cdf_from_pmf(pmf), "cdf_from_pmf")
    assert cdf[7] - cdf[6] == pmf[7]


# --------------------------------------------------------------------------
# Exercise 3 -- expectation and variance
# --------------------------------------------------------------------------


def test_3_expectation_of_sum_is_seven():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    got = attempt(lambda: dist.expectation_pmf(pmf), "expectation_pmf")
    assert got == 7


def test_3_variance_of_sum_is_35_over_6():
    pmf = attempt(dist.dice_sum_pmf, "dice_sum_pmf")
    got = attempt(lambda: dist.variance_pmf(pmf), "variance_pmf")
    assert got == Fraction(35, 6)


# --------------------------------------------------------------------------
# Exercise 4 -- linearity with a dependent pair
# --------------------------------------------------------------------------


def test_4_linearity_holds_for_the_dependent_pair():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    e_x = attempt(lambda: dist.expectation_over(outcomes, weight, D.first_die), "expectation_over")
    e_y = attempt(lambda: dist.expectation_over(outcomes, weight, D.dice_sum), "expectation_over")
    e_sum = attempt(
        lambda: dist.expectation_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o)),
        "expectation_over",
    )
    assert e_sum == e_x + e_y, "E[X + Y] must equal E[X] + E[Y] EXACTLY, even though X and Y are dependent"


# --------------------------------------------------------------------------
# Exercise 5 -- variance is not additive
# --------------------------------------------------------------------------


def test_5_variance_of_sum_is_not_the_naive_sum():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    var_x = attempt(lambda: dist.variance_over(outcomes, weight, D.first_die), "variance_over")
    var_y = attempt(lambda: dist.variance_over(outcomes, weight, D.dice_sum), "variance_over")
    var_sum = attempt(
        lambda: dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o)),
        "variance_over",
    )
    assert var_sum != var_x + var_y


def test_5_variance_of_sum_matches_the_full_covariance_formula():
    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
    var_x = attempt(lambda: dist.variance_over(outcomes, weight, D.first_die), "variance_over")
    var_y = attempt(lambda: dist.variance_over(outcomes, weight, D.dice_sum), "variance_over")
    cov_xy = attempt(
        lambda: dist.covariance_over(outcomes, weight, D.first_die, D.dice_sum), "covariance_over"
    )
    var_sum = attempt(
        lambda: dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o)),
        "variance_over",
    )
    assert var_sum == var_x + var_y + 2 * cov_xy


# --------------------------------------------------------------------------
# Exercise 6 -- Jensen's inequality
# --------------------------------------------------------------------------


def test_6_jensen_strict_for_a_die():
    outcomes, weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
    e_x2 = attempt(lambda: dist.expectation_over(outcomes, weight, lambda x: x * x), "expectation_over")
    e_x = attempt(lambda: dist.expectation_over(outcomes, weight, lambda x: x), "expectation_over")
    assert e_x2 > e_x**2


def test_6_jensen_gap_equals_variance():
    outcomes, weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
    e_x2 = attempt(lambda: dist.expectation_over(outcomes, weight, lambda x: x * x), "expectation_over")
    e_x = attempt(lambda: dist.expectation_over(outcomes, weight, lambda x: x), "expectation_over")
    var_x = attempt(lambda: dist.variance_over(outcomes, weight, lambda x: x), "variance_over")
    assert e_x2 - e_x**2 == var_x


# --------------------------------------------------------------------------
# Exercise 7 -- inverse-CDF discrete sampling
# --------------------------------------------------------------------------


def test_7_sampler_draws_only_pmf_values():
    pmf = {k: float(v) for k, v in D_pmf_for_starter()}
    rng = np.random.default_rng(0)
    draws = attempt(
        lambda: samp.sample_discrete_inverse_cdf(pmf, rng, 2_000), "sample_discrete_inverse_cdf"
    )
    assert set(draws.astype(int).tolist()) <= set(pmf)


def test_7_sampler_matches_pmf_within_tolerance():
    pmf = {k: float(v) for k, v in D_pmf_for_starter()}
    rng = np.random.default_rng(1)
    n = 100_000
    draws = attempt(
        lambda: samp.sample_discrete_inverse_cdf(pmf, rng, n), "sample_discrete_inverse_cdf"
    )
    values, counts = np.unique(draws, return_counts=True)
    empirical = dict(zip(values.astype(int), counts / n))
    for k, p in pmf.items():
        se = D.standard_error_of_proportion(p, n)
        close(empirical.get(k, 0.0), p, 4.0 * se, f"P(Y={k})")


def test_7_sampler_same_seed_is_reproducible():
    pmf = {k: float(v) for k, v in D_pmf_for_starter()}
    a = attempt(
        lambda: samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(7), 500),
        "sample_discrete_inverse_cdf",
    )
    b = attempt(
        lambda: samp.sample_discrete_inverse_cdf(pmf, np.random.default_rng(7), 500),
        "sample_discrete_inverse_cdf",
    )
    assert np.array_equal(a, b)


def D_pmf_for_starter():
    # A helper that does not itself depend on the reader's dice_sum_pmf,
    # so exercise 7's tests can run even before exercise 1 is solved.
    from itertools import product

    counts: dict[int, int] = {}
    for a, b in product(range(1, 7), range(1, 7)):
        total = a + b
        counts[total] = counts.get(total, 0) + 1
    return {k: Fraction(v, 36) for k, v in counts.items()}.items()


# --------------------------------------------------------------------------
# Exercise 8 -- exponential from scratch
# --------------------------------------------------------------------------


def test_8_exponential_scratch_is_nonnegative():
    rng = np.random.default_rng(3)
    draws = attempt(
        lambda: samp.sample_exponential_scratch(1.0, rng, 5_000), "sample_exponential_scratch"
    )
    assert (draws >= 0).all()


def test_8_exponential_scratch_mean_is_close_to_one_over_rate():
    rate = 2.0
    rng = np.random.default_rng(4)
    n = 100_000
    draws = attempt(
        lambda: samp.sample_exponential_scratch(rate, rng, n), "sample_exponential_scratch"
    )
    tol = 3.0 * D.standard_error_of_mean((1.0 / rate) ** 2, n)
    close(draws.mean(), 1.0 / rate, tol, "exponential sample mean")


def test_8_max_gap_of_identical_samples_is_zero():
    sample = np.array([1.0, 2.0, 3.0])
    got = attempt(lambda: samp.max_gap_statistic(sample, sample), "max_gap_statistic")
    assert got == 0.0


# --------------------------------------------------------------------------
# Exercise 9 -- Poisson as a Binomial limit
# --------------------------------------------------------------------------


def test_9_binomial_pmf_sums_to_one():
    total = 0.0
    for k in range(11):
        total += attempt(lambda k=k: dist.binomial_pmf(10, 0.3, k), "binomial_pmf")
    assert abs(total - 1.0) < 1e-9


def test_9_poisson_pmf_at_zero_matches_exp_of_minus_lambda():
    import math

    got = attempt(lambda: dist.poisson_pmf(2.0, 0), "poisson_pmf")
    assert abs(got - math.exp(-2.0)) < 1e-12


def test_9_gap_decreases_monotonically():
    lam = D.POISSON_LAMBDA
    gaps = [
        attempt(
            lambda n=n: dist.max_binomial_poisson_gap(n, lam / n, lam, D.POISSON_COMPARISON_KS),
            "max_binomial_poisson_gap",
        )
        for n in D.POISSON_LIMIT_NS
    ]
    assert all(a > b for a, b in zip(gaps, gaps[1:]))


# --------------------------------------------------------------------------
# Exercise 10 -- density above 1
# --------------------------------------------------------------------------


def test_10_uniform_density_is_two():
    got = attempt(lambda: dist.uniform_density(0.25, 0.0, 0.5), "uniform_density")
    assert got == 2.0


def test_10_numeric_integral_of_density_is_one():
    got = attempt(
        lambda: dist.numeric_integral(lambda x: dist.uniform_density(x, 0.0, 0.5), 0.0, 0.5, 50_000),
        "numeric_integral",
    )
    close(got, 1.0, 1e-6, "integral of the Uniform(0, 0.5) density")


# --------------------------------------------------------------------------
# The eighteen predictions
# --------------------------------------------------------------------------

EXPECTED: dict[str, object] = {
    "p_sum_seven": float(Fraction(1, 6)),
    "ratio_seven_to_two": 6,
    "cdf_at_twelve": 1.0,
    "cdf_difference_seven_six": float(Fraction(1, 6)),
    "expectation_of_sum": 7.0,
    "variance_of_sum": float(Fraction(35, 6)),
    "expectation_x_plus_y": float(Fraction(21, 2)),
    "linearity_holds_for_dependent_pair": True,
    "variance_naive_sum_holds": False,
    "variance_x_plus_y": float(Fraction(175, 12)),
    "jensen_strict_for_die": True,
    "jensen_gap_equals_variance": True,
    "discrete_sampler_reproducible": True,
    "exponential_scratch_nonnegative": True,
    "poisson_gap_decreases_monotonically": True,
    "uniform_density_value": 2.0,
    "uniform_density_exceeds_one": True,
    "uniform_integral_equals_one": True,
}

HINTS: dict[str, str] = {
    "ratio_seven_to_two": (
        "P(sum=7) = 6/36 and P(sum=2) = 1/36. Divide one by the other."
    ),
    "variance_naive_sum_holds": (
        "X and Y are dependent, and Var[X+Y] = Var[X] + Var[Y] + 2*Cov(X,Y). "
        "The covariance term here is not zero."
    ),
    "expectation_x_plus_y": (
        "Linearity holds regardless of dependence: E[X+Y] = E[X] + E[Y] = "
        "3.5 + 7."
    ),
}


@pytest.mark.parametrize("key", sorted(EXPECTED))
def test_predictions(key):
    got = need(answers.ANSWERS.get(key), f"answers.ANSWERS[{key!r}]")
    want = EXPECTED[key]
    hint = HINTS.get(key, "")
    if isinstance(want, bool) or isinstance(want, int):
        assert got == want, f"{key}: your answer {got!r}, expected {want!r}. {hint}"
    else:
        assert abs(float(got) - want) < 1e-6, (
            f"{key}: your answer {got!r}, expected {want!r}. {hint}"
        )


def test_every_answer_key_is_still_present():
    missing = sorted(set(EXPECTED) - set(answers.ANSWERS))
    assert not missing, f"answers.py is missing these keys: {missing}"

#!/usr/bin/env bash
# Tests for the Day 114 lab. Run from the lab directory:
#   bash tests/run_tests.sh
#
# The harness proves the lesson's claims by running code and reading real
# values, never by reading source:
#
#   * the two-dice sum pmf is exact Fraction arithmetic, and 7 is exactly
#     six times as likely as 2 -- the distribution is nowhere near uniform;
#   * the cdf is monotone, ends at exactly 1, and a cdf difference recovers
#     a pmf value exactly;
#   * expectation and variance computed from the definition agree with a
#     large seeded simulation, within three standard errors;
#   * E[X+Y] = E[X] + E[Y] EXACTLY for a dependent pair (X = first die,
#     Y = the sum), with no independence assumption anywhere;
#   * Var[X+Y] != Var[X] + Var[Y] for that same pair, but DOES equal
#     Var[X] + Var[Y] + 2*Cov(X,Y) exactly -- the asymmetry stated plainly;
#   * E[X^2] > (E[X])^2 for a die, and the gap is exactly Var[X];
#   * an inverse-CDF sampler written from scratch reproduces an arbitrary
#     discrete pmf within tolerance, and the same seed reproduces the same
#     draws;
#   * an exponential sampler written from scratch as -ln(U)/lambda agrees
#     with NumPy's own, both on sample mean and on a hand-written max-gap
#     statistic between their empirical cdfs (scipy is not installed);
#   * the Binomial(n, lambda/n) pmf converges to the Poisson(lambda) pmf as
#     n grows, measured as a monotonically shrinking maximum gap;
#   * Uniform(0, 0.5) has density 2 -- above 1 -- everywhere on its support,
#     while its numeric integral over that support is exactly 1;
#   * nothing is left behind on disk.
#
# Everything after the one-time install runs offline. Nothing binds a port,
# nothing writes outside the lab, nothing needs a key. Deterministic,
# non-interactive, exits 0 only if every check passes.
set -u

export PYTHONDONTWRITEBYTECODE=1

lab_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Bytecode left by an EARLIER command is not this run's litter. The README
# documents `pytest starter -q`, and running it writes .pyc files that would
# then fail the cleanliness check at the end of this script -- failing the
# reader for following the instructions. Clearing them here makes that final
# check measure what it claims to: what THIS run left behind. `.venv` is
# untouched, because the packages' own bytecode is theirs, not ours.
find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -exec rm -rf {} + 2>/dev/null || true

failures=0
checks=0

check() {
  local label="$1" ok="$2"
  checks=$((checks + 1))
  if [ "${ok}" = "yes" ]; then
    echo "  ok: ${label}"
  else
    echo "  FAIL: ${label}"
    failures=$((failures + 1))
  fi
}

check_eq() {
  # check_eq <label> <expected> <actual>
  if [ "$2" = "$3" ]; then
    check "$1" "yes"
  else
    check "$1 (expected [$2], got [$3])" "no"
  fi
}

# Resolve pytest: an explicit override, then this lab's .venv, then PATH.
# Fails loudly with instructions rather than silently skipping checks.
resolve_tool() {
  local tool="$1" override="$2"
  if [ -n "${override}" ] && [ -x "${override}" ]; then echo "${override}"; return 0; fi
  if [ -x "${lab_dir}/.venv/bin/${tool}" ]; then echo "${lab_dir}/.venv/bin/${tool}"; return 0; fi
  if command -v "${tool}" >/dev/null 2>&1; then command -v "${tool}"; return 0; fi
  return 1
}

pytest_bin="$(resolve_tool pytest "${PYTEST:-}")" || {
  echo "FAIL: pytest not found." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  echo "  Or point this suite at an existing pytest:" >&2
  echo "    PYTEST=/path/to/pytest bash tests/run_tests.sh" >&2
  exit 1
}

# The Python that owns that pytest is the one with numpy installed.
python_bin="$(dirname "${pytest_bin}")/python3"
if [ ! -x "${python_bin}" ]; then
  python_bin="$(command -v python3 || true)"
fi
if [ -z "${python_bin}" ]; then
  echo "FAIL: python3 not found on PATH." >&2
  exit 1
fi

if ! "${python_bin}" -c "import numpy" >/dev/null 2>&1; then
  echo "FAIL: numpy is not importable from ${python_bin}." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  exit 1
fi

echo "Day 114 — Random Variables and Distributions"
echo

# --------------------------------------------------------------------------
echo "1. The tools and the versions this lab was written against"
# --------------------------------------------------------------------------

versions="$("${python_bin}" - <<'PY'
import platform
import sys
from importlib.metadata import version

print(f"python   {platform.python_version()}")
for name in ("numpy", "pytest"):
    print(f"{name:<8} {version(name)}")
print(f"platform {platform.platform()}")
print(f"exe      {sys.executable.rsplit('/', 3)[-1]}")
PY
)"
echo "${versions}" | sed 's/^/  /'

pinned_numpy="$(grep -E '^numpy==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
installed_numpy="$("${python_bin}" -c "from importlib.metadata import version; print(version('numpy'))")"
check_eq "installed numpy matches requirements.txt" "${pinned_numpy}" "${installed_numpy}"

major="$("${python_bin}" -c "import numpy; print(numpy.__version__.split('.')[0])")"
check_eq "numpy is version 2 or later" "2" "${major}"

# --------------------------------------------------------------------------
echo
echo "2. Every reference script runs and every assertion inside it holds"
# --------------------------------------------------------------------------

for script in 01_pmf_of_a_sum 02_cdf_from_pmf 03_expectation_and_variance \
              04_linearity_with_dependence 05_variance_is_not_additive \
              06_jensens_inequality 07_inverse_cdf_discrete_sampler \
              08_exponential_from_scratch 09_poisson_as_binomial_limit \
              10_density_above_one; do
  out="$(cd "${lab_dir}/examples" && "${python_bin}" "${script}.py" 2>&1)"
  status=$?
  if [ "${status}" -ne 0 ]; then
    check "${script}.py exits 0" "no"
    echo "${out}" | tail -5 | sed 's/^/      /'
  else
    check "${script}.py exits 0" "yes"
  fi
  case "${out}" in
    *"${script}.py: every assertion held."*)
      check "${script}.py reports every assertion held" "yes" ;;
    *) check "${script}.py reports every assertion held" "no" ;;
  esac
done

# --------------------------------------------------------------------------
echo
echo "3. The reference pytest suite: real values, real exceptions"
# --------------------------------------------------------------------------

ref_out="$(cd "${lab_dir}" && "${pytest_bin}" examples -q -p no:cacheprovider 2>&1)"
ref_status=$?
echo "${ref_out}" | tail -3 | sed 's/^/  /'
if [ "${ref_status}" -eq 0 ]; then
  check "pytest examples exits 0" "yes"
else
  check "pytest examples exits 0" "no"
fi
case "${ref_out}" in
  *" failed"*) check "no test in the reference suite failed" "no" ;;
  *)           check "no test in the reference suite failed" "yes" ;;
esac
ref_passed="$(printf '%s\n' "${ref_out}" | grep -o '[0-9][0-9]* passed' | head -1 | cut -d' ' -f1)"
if [ "${ref_passed:-0}" -ge 60 ]; then
  check "the reference suite ran at least 60 tests (ran ${ref_passed})" "yes"
else
  check "the reference suite ran at least 60 tests (ran ${ref_passed:-0})" "no"
fi

# --------------------------------------------------------------------------
echo
echo "4. The starter suite skips unattempted work instead of failing it"
# --------------------------------------------------------------------------

start_out="$(cd "${lab_dir}" && "${pytest_bin}" starter -q -p no:cacheprovider 2>&1)"
start_status=$?
echo "${start_out}" | tail -3 | sed 's/^/  /'
if [ "${start_status}" -eq 0 ]; then
  check "pytest starter exits 0 on an untouched checkout" "yes"
else
  check "pytest starter exits 0 on an untouched checkout" "no"
fi
case "${start_out}" in
  *" failed"*) check "the starter suite reports no failures" "no" ;;
  *)           check "the starter suite reports no failures" "yes" ;;
esac
case "${start_out}" in
  *skipped*) check "unwritten exercises are reported as skipped, not passed" "yes" ;;
  *) check "unwritten exercises are reported as skipped, not passed" "no" ;;
esac

# The import guard. Both directories contain modules called `distributions`,
# `sampling`, `dataset` and `answers`, and pytest imports test files by
# putting their directory on sys.path -- so collecting both suites at once
# would otherwise let the starter tests import the REFERENCE solution and
# report unwritten exercises as passing. Each directory's conftest.py
# prevents that. This check proves it still does: across both suites, the
# skip count must be unchanged.
both_out="$(cd "${lab_dir}" && "${pytest_bin}" -q -p no:cacheprovider 2>&1)"
start_skipped="$(printf '%s\n' "${start_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
both_skipped="$(printf '%s\n' "${both_out}" | grep -o '[0-9][0-9]* skipped' | head -1 | cut -d' ' -f1)"
check_eq "collecting both suites at once does not turn skips into passes" \
  "${start_skipped:-none}" "${both_skipped:-none}"

# --------------------------------------------------------------------------
echo
echo "5. The lesson's claims, checked one value at a time"
# --------------------------------------------------------------------------

facts="$(cd "${lab_dir}/examples" && "${python_bin}" - <<'PY'
from fractions import Fraction

import numpy as np

import dataset as D
import distributions as dist
import sampling as samp

# -- exercise 1: the pmf of a sum -------------------------------------------
pmf = dist.dice_sum_pmf()
print("pmf_p7", pmf[7])
print("pmf_p2", pmf[2])
print("pmf_ratio", pmf[7] / pmf[2])
print("pmf_not_uniform", len(set(pmf.values())) > 1)

# -- exercise 2: the cdf ------------------------------------------------------
cdf = dist.cdf_from_pmf(pmf)
values = [cdf[k] for k in sorted(cdf)]
print("cdf_monotone", all(a <= b for a, b in zip(values, values[1:])))
print("cdf_ends_at_one", cdf[12] == 1)
print("cdf_diff_matches_pmf", cdf[7] - cdf[6] == pmf[7])

# -- exercise 3: expectation and variance ------------------------------------
exact_mean = dist.expectation_pmf(pmf)
exact_var = dist.variance_pmf(pmf)
print("exact_mean", exact_mean)
print("exact_var", exact_var)

rng = np.random.default_rng(D.SEED)
first = rng.integers(1, 7, size=D.EV_SIMULATION_TRIALS)
second = rng.integers(1, 7, size=D.EV_SIMULATION_TRIALS)
sample = (first + second).astype(float)
mean_tol = 3.0 * D.standard_error_of_mean(float(exact_var), D.EV_SIMULATION_TRIALS)
print("mean_within_tol", abs(sample.mean() - float(exact_mean)) < mean_tol)

# -- exercises 4, 5: linearity and non-additivity ----------------------------
outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT
E_X = dist.expectation_over(outcomes, weight, D.first_die)
E_Y = dist.expectation_over(outcomes, weight, D.dice_sum)
E_XY = dist.expectation_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))
print("E_X", E_X)
print("E_Y", E_Y)
print("linearity_exact", E_XY == E_X + E_Y)

Var_X = dist.variance_over(outcomes, weight, D.first_die)
Var_Y = dist.variance_over(outcomes, weight, D.dice_sum)
Var_XY = dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))
Cov_XY = dist.covariance_over(outcomes, weight, D.first_die, D.dice_sum)
print("variance_naive_wrong", Var_XY != Var_X + Var_Y)
print("variance_full_formula_exact", Var_XY == Var_X + Var_Y + 2 * Cov_XY)
print("covariance_nonzero", Cov_XY != 0)

# -- exercise 6: Jensens inequality -------------------------------------------
die_outcomes, die_weight = D.DIE_FACES, D.ONE_DIE_WEIGHT
EX2 = dist.expectation_over(die_outcomes, die_weight, lambda x: x * x)
EX = dist.expectation_over(die_outcomes, die_weight, lambda x: x)
VarX = dist.variance_over(die_outcomes, die_weight, lambda x: x)
print("jensen_strict", EX2 > EX**2)
print("jensen_gap_exact", (EX2 - EX**2) == VarX)

# -- exercise 7: inverse-CDF discrete sampler --------------------------------
pmf_float = {k: float(v) for k, v in pmf.items()}
rng7 = np.random.default_rng(D.SEED)
draws = samp.sample_discrete_inverse_cdf(pmf_float, rng7, D.DISCRETE_SAMPLER_TRIALS)
values7, counts7 = np.unique(draws, return_counts=True)
empirical = dict(zip(values7.astype(int), counts7 / D.DISCRETE_SAMPLER_TRIALS))
worst_se = max(D.standard_error_of_proportion(p, D.DISCRETE_SAMPLER_TRIALS) for p in pmf_float.values())
max_gap7 = max(abs(pmf_float[k] - empirical.get(k, 0.0)) for k in pmf_float)
print("sampler_within_tol", max_gap7 < 3.0 * worst_se)

rng_a = np.random.default_rng(D.SEED)
rng_b = np.random.default_rng(D.SEED)
d_a = samp.sample_discrete_inverse_cdf(pmf_float, rng_a, 1000)
d_b = samp.sample_discrete_inverse_cdf(pmf_float, rng_b, 1000)
print("sampler_reproducible", np.array_equal(d_a, d_b))

# -- exercise 8: exponential from scratch, and the max-gap statistic --------
rng8 = np.random.default_rng(D.SEED)
scratch = samp.sample_exponential_scratch(D.EXPONENTIAL_RATE, rng8, D.EXPONENTIAL_SAMPLE_SIZE)
built_in = rng8.exponential(scale=1.0 / D.EXPONENTIAL_RATE, size=D.EXPONENTIAL_SAMPLE_SIZE)
target = 1.0 / D.EXPONENTIAL_RATE
tol8 = 3.0 * D.standard_error_of_mean(target**2, D.EXPONENTIAL_SAMPLE_SIZE)
print("exp_scratch_mean_ok", abs(scratch.mean() - target) < tol8)
print("exp_builtin_mean_ok", abs(built_in.mean() - target) < tol8)
gap8 = samp.max_gap_statistic(scratch, built_in)
threshold8 = D.dkw_two_sample_threshold(D.EXPONENTIAL_SAMPLE_SIZE, D.EXPONENTIAL_SAMPLE_SIZE)
print("exp_max_gap_below_threshold", gap8 < threshold8)
print("exp_max_gap", f"{gap8:.6f}")
print("exp_threshold", f"{threshold8:.6f}")

# -- exercise 9: Poisson as a Binomial limit ---------------------------------
lam = D.POISSON_LAMBDA
gaps9 = [
    dist.max_binomial_poisson_gap(n, lam / n, lam, D.POISSON_COMPARISON_KS)
    for n in D.POISSON_LIMIT_NS
]
print("poisson_limit_monotone", all(a > b for a, b in zip(gaps9, gaps9[1:])))
print("poisson_limit_last_tiny", gaps9[-1] < 1e-3)

# -- exercise 10: density above 1 ---------------------------------------------
density10 = dist.uniform_density(0.25, D.UNIFORM_LOW, D.UNIFORM_HIGH)
integral10 = dist.numeric_integral(
    lambda x: dist.uniform_density(x, D.UNIFORM_LOW, D.UNIFORM_HIGH),
    D.UNIFORM_LOW, D.UNIFORM_HIGH, 100_000,
)
print("density_is_two", density10 == 2.0)
print("density_exceeds_one", density10 > 1.0)
print("integral_is_one", round(integral10, 6) == 1.0)
PY
)"

get() { printf '%s\n' "${facts}" | grep "^$1 " | cut -d' ' -f2-; }

check_eq "P(sum=7) is exactly 1/6" "1/6" "$(get pmf_p7)"
check_eq "P(sum=2) is exactly 1/36" "1/36" "$(get pmf_p2)"
check_eq "7 is exactly six times as likely as 2" "6" "$(get pmf_ratio)"
check_eq "the distribution is not uniform" "True" "$(get pmf_not_uniform)"
check_eq "the cdf is monotone non-decreasing" "True" "$(get cdf_monotone)"
check_eq "the cdf ends at exactly 1" "True" "$(get cdf_ends_at_one)"
check_eq "F(7) - F(6) equals P(sum=7) exactly" "True" "$(get cdf_diff_matches_pmf)"
check_eq "E[Y] is exactly 7" "7" "$(get exact_mean)"
check_eq "Var[Y] is exactly 35/6" "35/6" "$(get exact_var)"
check_eq "the simulated mean lands within 3 standard errors" "True" "$(get mean_within_tol)"
check_eq "E[X] is 7/2" "7/2" "$(get E_X)"
check_eq "E[Y] (joint) is 7" "7" "$(get E_Y)"
check_eq "E[X+Y] equals E[X]+E[Y] EXACTLY for the dependent pair" "True" "$(get linearity_exact)"
check_eq "Var[X+Y] does NOT equal Var[X]+Var[Y]" "True" "$(get variance_naive_wrong)"
check_eq "Var[X+Y] EXACTLY equals Var[X]+Var[Y]+2*Cov(X,Y)" "True" "$(get variance_full_formula_exact)"
check_eq "Cov(X,Y) is non-zero for this dependent pair" "True" "$(get covariance_nonzero)"
check_eq "E[X^2] > (E[X])^2 for a die (Jensen)" "True" "$(get jensen_strict)"
check_eq "the Jensen gap equals Var[X] exactly" "True" "$(get jensen_gap_exact)"
check_eq "the inverse-CDF sampler matches the pmf within tolerance" "True" "$(get sampler_within_tol)"
check_eq "the same seed reproduces identical discrete draws" "True" "$(get sampler_reproducible)"
check_eq "the from-scratch exponential sampler's mean is within tolerance" "True" "$(get exp_scratch_mean_ok)"
check_eq "NumPy's own exponential sampler's mean is within tolerance" "True" "$(get exp_builtin_mean_ok)"
check_eq "the max-gap statistic is below the DKW-derived threshold" "True" "$(get exp_max_gap_below_threshold)"
echo "  (measured on this run: max-gap statistic $(get exp_max_gap) against threshold $(get exp_threshold) -- reported, not asserted to a value)"
check_eq "the Binomial-to-Poisson pmf gap shrinks monotonically with n" "True" "$(get poisson_limit_monotone)"
check_eq "the gap at n=10,000 is under 0.001" "True" "$(get poisson_limit_last_tiny)"
check_eq "Uniform(0, 0.5)'s density is exactly 2" "True" "$(get density_is_two)"
check_eq "that density exceeds 1" "True" "$(get density_exceeds_one)"
check_eq "the numeric integral of that density is 1 to six decimals" "True" "$(get integral_is_one)"

# --------------------------------------------------------------------------
echo
echo "6. The harness can actually fail"
# --------------------------------------------------------------------------

# A green test suite proves nothing until you have watched it go red. This
# section re-runs the reference pytest suite with ONE assertion deliberately
# swapped for a wrong one -- claiming Var[X+Y] == Var[X]+Var[Y], the naive
# belief the lesson spends a whole exercise refuting -- and asserts that the
# re-run reports exactly one failure and a non-zero exit. If this section
# passes, section 5 is not decorative.

self_test_marker="test_variance_of_sum_is_not_the_naive_sum"
original_file="${lab_dir}/examples/test_reference.py"
backup_file="$(mktemp)"
cp "${original_file}" "${backup_file}"

python3 - "${original_file}" "${self_test_marker}" <<'PY'
import re
import sys

path, marker = sys.argv[1], sys.argv[2]
text = open(path).read()
needle = (
    "def test_variance_of_sum_is_not_the_naive_sum():\n"
    "    outcomes, weight = D.TWO_DICE_SPACE, D.TWO_DICE_WEIGHT\n"
    "    var_x = dist.variance_over(outcomes, weight, D.first_die)\n"
    "    var_y = dist.variance_over(outcomes, weight, D.dice_sum)\n"
    "    var_sum = dist.variance_over(outcomes, weight, lambda o: D.first_die(o) + D.dice_sum(o))\n"
    "    assert var_sum != var_x + var_y\n"
)
replacement = needle.replace("assert var_sum != var_x + var_y", "assert var_sum == var_x + var_y")
assert needle in text, "self-test marker not found -- test_reference.py has drifted"
text = text.replace(needle, replacement)
open(path, "w").write(text)
PY

self_out="$(cd "${lab_dir}" && "${pytest_bin}" examples -q -p no:cacheprovider 2>&1)"
self_status=$?
cp "${backup_file}" "${original_file}"
rm -f "${backup_file}"
find "${lab_dir}/examples" -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true

if [ "${self_status}" -ne 0 ]; then
  check "a deliberately wrong assertion makes the reference suite exit non-zero (${self_status})" "yes"
else
  check "a deliberately wrong assertion makes the reference suite exit non-zero" "no"
fi
case "${self_out}" in
  *"${self_test_marker}"*)
    check "the failing test is named in the output" "yes" ;;
  *) check "the failing test is named in the output" "no" ;;
esac
case "${self_out}" in
  *"1 failed"*)
    check "the summary line counts exactly one failure" "yes" ;;
  *) check "the summary line counts exactly one failure" "no" ;;
esac

# --------------------------------------------------------------------------
echo
echo "7. Nothing was left behind"
# --------------------------------------------------------------------------

# `.venv` is pruned from both searches below. The virtual environment ships
# NumPy's and pytest's own precompiled bytecode -- hundreds of __pycache__
# directories that came with the packages and have nothing to do with whether
# THIS lab tidied up after itself. Searching them would report a failure the
# reader cannot fix and did not cause. Everything the lab itself writes lives
# outside `.venv`, which is exactly what these two checks look at.

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -print -quit 2>/dev/null | grep -q .; then
  check "no __pycache__ directory left by the lab's own code" "no"
else
  check "no __pycache__ directory left by the lab's own code" "yes"
fi

if find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -print -quit 2>/dev/null | grep -q .; then
  check "no .pytest_cache directory left under the lab" "no"
else
  check "no .pytest_cache directory left under the lab" "yes"
fi

if grep -rqE 'urlopen|requests\.|socket\.|http://|https://' \
     "${lab_dir}/examples" "${lab_dir}/starter" 2>/dev/null; then
  check "no lab source opens a network connection" "no"
else
  check "no lab source opens a network connection" "yes"
fi

echo
echo "${checks} checks, ${failures} failure(s)."
[ "${failures}" -eq 0 ]