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

import math
import numpy as np
from scipy.stats import binomtest

rng = np.random.default_rng(7)
few = np.r_[np.ones(10), np.zeros(990)]
many = np.r_[np.ones(105), -np.ones(95), np.zeros(800)]
assert np.isclose(few.mean(), many.mean())
p_few = binomtest(10, 10, p=0.5).pvalue
p_many = binomtest(105, 200, p=0.5).pvalue
assert p_few < 0.01 and p_many > 0.05
resampled = rng.choice(many, size=(2000, len(many)), replace=True).mean(1)
interval = np.quantile(resampled, [0.025, 0.975])
assert interval[0] < 0 < interval[1]

def pass_at_k(n, c, k):
    if not (0 <= c <= n and 1 <= k <= n):
        raise ValueError("require 0 <= c <= n and 1 <= k <= n")
    if n - c < k:
        return 1.0
    return -math.expm1(sum(math.log1p(-c / (n - i)) for i in range(k)))

assert math.isclose(pass_at_k(10, 2, 3), 1 - math.comb(8, 3) / math.comb(10, 3))
assert pass_at_k(10, 0, 3) == 0 and pass_at_k(10, 10, 3) == 1
print("paired p-values:", p_few, p_many, "bootstrap interval:", interval)