# Source: content/notes/math/statistics.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from scipy import stats
from sklearn.metrics import f1_score

rng = np.random.default_rng(51)
d = np.array([1., 2., 0., 1., 1.])
result = stats.ttest_1samp(d, 0.)
se = stats.sem(d)
interval = stats.t.interval(.95, len(d)-1, loc=d.mean(), scale=se)
assert .033 < result.pvalue < .035
assert interval[0] > 0
assert stats.binomtest(10, 10, .5).pvalue == 2/1024

# Known-sigma coverage: simulation checks implementation, not a theorem.
samples = rng.normal(2., 3., size=(5000, 36))
means = samples.mean(1)
coverage = np.mean(np.abs(means-2.) <= stats.norm.ppf(.975)*3/6)
assert .93 < coverage < .97

y = rng.binomial(1, .2, 500)
a = np.where(rng.random(500) < .12, 1-y, y)
b = np.where(rng.random(500) < .18, 1-y, y)
differences = []
for _ in range(1500):
    idx = rng.integers(0, len(y), len(y))
    # zero_division=0 explicitly defines a degenerate resample's F1.
    differences.append(f1_score(y[idx], a[idx], zero_division=0) -
                       f1_score(y[idx], b[idx], zero_division=0))
lo, hi = np.quantile(differences, [.025, .975])
assert np.isfinite([lo, hi]).all() and lo <= hi
print("paired t:", result.pvalue, interval)
print("normal interval coverage:", coverage, "paired F1 interval:", (lo, hi))