# Independent re-implementation of the six-order N/s/r decomposition, plus a common-census-base sensitivity
# using Census's 2010-2020 intercensal July 1, 2018 resident estimates (consistent with the 2020 census).
import csv, json, sys
from fractions import Fraction as Q
from itertools import permutations

r1 = json.load(open(sys.argv[1]))
cells = r1["endpoint_age_cells"]
bins = [(0,0),(1,4),(5,14),(15,24),(25,34),(35,44),(45,54),(55,64),(65,74),(75,84),(85,200)]

def decompose(d0, n0, d1, n1):
    def f(d, n):
        N = sum(n); return (N, [Q(x, N) for x in n], [Q(a, b) for a, b in zip(d, n)])
    def D(N, s, r): return N * sum(a*b for a, b in zip(s, r))
    start, end = f(d0, n0), f(d1, n1)
    total = [Q(0)]*3
    for order in permutations(range(3)):
        cur = list(start); before = D(*cur)
        for k in order:
            cur[k] = end[k]; after = D(*cur); total[k] += after - before; before = after
    return [float(t/6) for t in total]

d0 = [c["baseline_deaths"] for c in cells]; n0 = [c["baseline_population"] for c in cells]
d1 = [c["endpoint_deaths"] for c in cells]; n1 = [c["endpoint_population"] for c in cells]
print("as published (WONDER denominators):", [round(x, 1) for x in decompose(d0, n0, d1, n1)])

pop = {}
for row in csv.DictReader(open(sys.argv[2])):
    if row["SEX"] == "0" and row["AGE"] != "999":
        pop[int(row["AGE"])] = int(row["POPESTIMATE2018"])
n0i = [sum(v for a, v in pop.items() if lo <= a <= hi) for lo, hi in bins]
print("2018 bins, WONDER vs intercensal:")
for (lo, hi), a, b in zip(bins, n0, n0i):
    print(f"  {lo}-{hi}: {a} {b} ({100*(b/a-1):+.2f}%)")
print("totals", sum(n0), sum(n0i), f"{100*(sum(n0i)/sum(n0)-1):+.3f}%")
print("2018 intercensal denominators:", [round(x, 1) for x in decompose(d0, n0i, d1, n1)])