"""How the sealed rule behaves under assumed truths, at the n actually reached.

    python research/so-accepted-answer/branches.py > branches.txt

Written AFTER the data, and it is honest anyway, for one reason that must be
checked before anybody reuses it: nothing here reads the observed share. The
inputs are the assumed true share, the sample size reached, and the bar from
rule.py. A branch table that resamples around the value that came back makes
every distant bar look dead ([[branch-odds-before-the-data]]); this one draws
from the assumed truth, so it answers "if the world were like that, what would
this rule have said", which is the only question a branch table can answer.

The one thing the data supplied is n — 2,026 for the old cohort and 893 for
the new one, because the new cohort spent the day's API quota before it
reached the target of 2,000. That is a weaker rule than the one sealed, and
this table is where that shows.
"""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import rule  # noqa: E402

from math import comb

TRUTHS = [0.60, 0.55, 0.51, 0.50, 0.45, 0.30, 0.16]
NS = [("old", 2026), ("new", 893), ("sealed target", 2000)]


def outcome_probs(p, n):
    """Exact binomial sum over every k, through the rule's own verdict."""
    out = {"survived": 0.0, "killed": 0.0, "inconclusive": 0.0}
    # log-space binomial pmf, because comb(2026, k) is enormous
    from math import lgamma, log, exp
    lp, lq = log(p), log(1 - p)
    lnc = lambda k: (lgamma(n + 1) - lgamma(k + 1) - lgamma(n - k + 1))
    for k in range(n + 1):
        w = exp(lnc(k) + k * lp + (n - k) * lq)
        if w < 1e-15:
            continue
        out[rule.verdict(k, n)["verdict"]] += w
    return out


def main():
    print("The rule under assumed truths. Bar: survives if the Wilson lower")
    print("bound is above %.2f, killed if the upper bound is below it." % rule.BAR)
    print("Computed after the data, from the assumed truth and n only.")
    print()
    print("%-14s %5s  %9s %9s %9s" % ("n", "true f", "survived", "killed", "inconcl."))
    for label, n in NS:
        for p in TRUTHS:
            o = outcome_probs(p, n)
            print("%-14s %5.2f  %9.3f %9.3f %9.3f"
                  % ("%s (%d)" % (label, n), p, o["survived"], o["killed"],
                     o["inconclusive"]))
        print()
    assert comb  # imported for the reader; the sum above is done in log space


if __name__ == "__main__":
    main()