#!/usr/bin/env python3
"""Exact hypothetical probabilities, not a model benchmark.

A: every fixed task succeeds with probability 1/2 on each independent run.
B: half the fixed task types always succeed; half always fail.
Task types have equal weight. Repeated runs keep the task type fixed.
Uses the Python standard library; no model calls, randomness, or network.
"""

from fractions import Fraction
from itertools import product
from math import comb


def metrics(task_probabilities, k):
    totals = {"single": Fraction(0), "any": Fraction(0), "all": Fraction(0)}
    for p in task_probabilities:
        mass = Fraction(0)
        for outcomes in product((False, True), repeat=k):
            weight = Fraction(1)
            for passed in outcomes:
                weight *= p if passed else 1 - p
            mass += weight
            totals["single"] += weight * outcomes[0]
            totals["any"] += weight * any(outcomes)
            totals["all"] += weight * all(outcomes)
        assert mass == 1
    return {key: value / len(task_probabilities) for key, value in totals.items()}


def observed_estimates(n, c, k):
    """Per-task estimates from an observed pool; n >= k >= 1, 0 <= c <= n."""
    if not 1 <= k <= n or not 0 <= c <= n:
        raise ValueError("Require 1 <= k <= n and 0 <= c <= n")
    return 1 - Fraction(comb(n - c, k), comb(n, k)), Fraction(comb(c, k), comb(n, k))


def main():
    populations = {"A": (Fraction(1, 2),), "B": (Fraction(1), Fraction(0))}
    expected = {
        "A": {"single": Fraction(1, 2), "any": Fraction(15, 16), "all": Fraction(1, 16)},
        "B": {"single": Fraction(1, 2), "any": Fraction(1, 2), "all": Fraction(1, 2)},
    }
    print("Hypothetical exact probabilities; four independent runs per fixed task")
    print("Agent  First chosen  Any of four  All four")
    for name, probabilities in populations.items():
        result = metrics(probabilities, 4)
        assert result == expected[name]
        assert result["any"] == sum(1 - (1 - p) ** 4 for p in probabilities) / len(probabilities)
        assert result["all"] == sum(p ** 4 for p in probabilities) / len(probabilities)
        print(f"{name:5}  {float(result['single']):11.2%}  {float(result['any']):11.2%}  {float(result['all']):8.2%}")

    # Verify the finite-pool estimators against enumerated subsets for every
    # possible observed binary pool up to six runs, at every supported k.
    from itertools import combinations

    for n in range(1, 7):
        for outcomes in product((False, True), repeat=n):
            for k in range(1, n + 1):
                subsets = tuple(combinations(outcomes, k))
                exact = (Fraction(sum(map(any, subsets)), len(subsets)),
                         Fraction(sum(map(all, subsets)), len(subsets)))
                assert observed_estimates(n, sum(outcomes), k) == exact
    print("Exact enumeration and observed-pool checks passed.")


if __name__ == "__main__":
    main()