"""How big a link do the new years still allow? (Exploratory; not part of the frozen test.)

For each block written by outsample_test.py: the ratio of observed to expected overlaps with an approximate 95%
interval, from the exact Poisson interval for the observed count divided by the expected count; and the chance of
so few overlaps if every date fell on an earthquake day at the paper's 2010-2012 rate, 17 of 48 (binomial).
The two blocks are also added together, as a rough summary (they were tested separately).

Standard library only. Usage: python3 -I intervals.py <outsample json> [<outsample json> ...] [--out file.json]
"""
import json, math, sys

RATE = 17 / 48


def pois_cdf(k, mu):
    if k < 0:
        return 0.0
    term = math.exp(-mu); total = term
    for i in range(1, k + 1):
        term *= mu / i; total += term
    return total


def solve(f, lo, hi):
    for _ in range(200):
        mid = (lo + hi) / 2
        if f(mid) > 0:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2


def poisson_ci(k, conf=0.95):
    a = (1 - conf) / 2
    # lower: P(X >= k | mu) = a; upper: P(X <= k | mu) = a. solve() moves up while f > 0, i.e. below the root.
    lower = 0.0 if k == 0 else solve(lambda mu: a - (1 - pois_cdf(k - 1, mu)), 1e-9, 10 * k + 10)
    upper = solve(lambda mu: pois_cdf(k, mu) - a, 1e-9, 10 * k + 50)
    return lower, upper


def binom_cdf(k, n, p):
    return sum(math.comb(n, i) * p**i * (1 - p)**(n - i) for i in range(0, k + 1))


def summarize(name, k, e, m):
    lo, hi = poisson_ci(k)
    return {"block": name, "observed": k, "expected": round(e, 3), "ratio": round(k / e, 3),
            "ratio_95_interval_approx": [round(lo / e, 2), round(hi / e, 2)],
            "expected_at_paper_rate": round(m * RATE, 1), "dates": m,
            "P_at_most_observed_at_paper_rate": float(f"{binom_cdf(k, m, RATE):.2e}")}


def main():
    args = sys.argv[1:]; out = None
    if "--out" in args:
        i = args.index("--out"); out = args[i + 1]; del args[i:i + 2]
    blocks = [json.load(open(p, encoding="utf-8")) for p in args]
    res = [summarize(b["years"], b["observed"], b["expected_exact"], b["m_dates"]) for b in blocks]
    if len(blocks) > 1:
        res.append(summarize(" + ".join(b["years"] for b in blocks), sum(b["observed"] for b in blocks),
                             sum(b["expected_exact"] for b in blocks), sum(b["m_dates"] for b in blocks)))
    print(json.dumps(res, indent=1))
    if out:
        with open(out, "w", encoding="utf-8", newline="\n") as f:
            json.dump(res, f, indent=1); f.write("\n")


if __name__ == "__main__":
    main()