"""Export model4 as a table of 'cells' so the web page can re-weight the discrete priors exactly, with no Monte Carlo in the browser.
A cell is one combination of the discrete prior choices; for each cell we store its base prior mass and, for each column, the
likelihood-weighted mass in each class and in the headline splits. Changing a prior marginal just rescales cells. -> cells4.json"""
import os
os.environ.setdefault("OMP_NUM_THREADS", "1"); os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
import json
import sys
from multiprocessing import Pool
import numpy as np
import model4 as M

HERE = os.path.dirname(os.path.abspath(__file__))
DIMS = ["fast", "m", "barrier", "regime", "cap", "hazard"]
SHAPE = (2, 5, 2, 3, 2, 2)
EXTRA = ["alone", "alone_sterile", "stuck_barrier", "other_now"]
NC = len(M.CLASSES)
HEDGES = np.arange(-64, 12.5, 2.0)       # log10 of the expected number of civilisations ever in the Galaxy


def job(args):
    seed, n, tilt = args
    r = M.run(M.vary(M.BASE, n=n, seed=seed, tilt=tilt), keep=True)
    s = r["_s"]
    idx = np.ravel_multi_index((s["fast"].astype(int), np.minimum(s["m"], 4), s["barrier"].astype(int), s["regime"], s["cap_short"].astype(int), s["hazard"].astype(int)), SHAPE)
    ncell = int(np.prod(SHAPE))
    out = dict(prior=np.bincount(idx, weights=s["imp"], minlength=ncell), nprior=float(s["imp"].sum()), schools={}, hist={})
    ln = np.log10(np.maximum(s["N_ever"], 1e-300)).clip(HEDGES[0], HEDGES[-1] - 1e-9)
    lnp = np.log10(np.maximum(s["prior_N_ever"], 1e-300)).clip(HEDGES[0], HEDGES[-1] - 1e-9)     # the prior uses the unconditioned gate delay
    out["hist"]["prior"] = np.histogram(lnp, HEDGES, weights=s["imp"])[0]
    alone = s["N_ever"] < 1
    p_now = -np.expm1(-np.minimum(s["N_now"], 700))
    for sc, w in s["weights"].items():
        mass = np.zeros((ncell, NC + len(EXTRA)))
        for c in range(NC):
            mk = s["cls"] == c
            mass[:, c] = np.bincount(idx[mk], weights=w[mk], minlength=ncell)
        mass[:, NC] = np.bincount(idx[alone], weights=w[alone], minlength=ncell)
        mk = s["sub"]["alone_sterile"]; mass[:, NC + 1] = np.bincount(idx[mk], weights=w[mk], minlength=ncell)
        mk = s["sub"]["stuck_barrier"]; mass[:, NC + 2] = np.bincount(idx[mk], weights=w[mk], minlength=ncell)
        mass[:, NC + 3] = np.bincount(idx, weights=w * p_now, minlength=ncell)
        out["schools"][sc] = mass
        out["hist"][sc] = np.histogram(ln, HEDGES, weights=w)[0]
    return tilt, out


def main(n=400_000, reps=8):
    with Pool(int(os.environ.get("FERMI_PROCS", "4"))) as pool:
        res = pool.map(job, [(300 + i, n, t) for i in range(reps) for t in (False, True)], chunksize=1)
    plain = [o for t, o in res if not t]; tilted = [o for t, o in res if t]
    prior = sum(o["prior"] for o in plain) / sum(o["nprior"] for o in plain)
    B = M.BASE
    out = dict(dims=DIMS, shape=list(SHAPE), classes=M.CLASSES, extra=EXTRA, prior=[round(float(x), 7) for x in prior], schools={})
    for sc in M.SCHOOLS:
        src = tilted if sc == "sia" else plain
        mass = sum(o["schools"][sc] for o in src)
        tot = mass[:, :NC].sum()
        if sc == "sia":      # tilted runs sample cells with different frequencies; that is already inside the weights
            pass
        out["schools"][sc] = [[float(f"{v / tot:.4g}") for v in row] for row in mass]
    hp = sum(o["hist"]["prior"] for o in plain)
    out["N_ever_hist"] = dict(edges=HEDGES.tolist(), prior=[round(float(v), 5) for v in hp / hp.sum()])
    for sc in M.SCHOOLS:
        hh = sum(o["hist"][sc] for o in (tilted if sc == "sia" else plain))
        out["N_ever_hist"][sc] = [round(float(v), 5) for v in hh / hh.sum()]
    mw = B["m_weights"]
    out["base_marginals"] = dict(fast=[1 - B["w_fast"], B["w_fast"]], m=[mw.get(k, 0) for k in range(5)], barrier=[1 - B["q_barrier"], B["q_barrier"]],
                                 regime=[1 - B["p_stealth"] - B["q_cc"], B["q_cc"], B["p_stealth"]], cap=[1 - B["p_cap_short"], B["p_cap_short"]],
                                 hazard=[1 - B["p_persist_hazard"], B["p_persist_hazard"]])
    json.dump(out, open(os.path.join(HERE, "cells4.json"), "w"), separators=(",", ":"))
    print("cells", len(prior), "non-empty", int((prior > 0).sum()), "file KB", os.path.getsize(os.path.join(HERE, "cells4.json")) // 1024)
    for sc in out["schools"]:
        a = np.array(out["schools"][sc]); print(sc, np.round(a[:, :NC].sum(0), 3), "alone", round(a[:, NC].sum(), 3))


if __name__ == "__main__":
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 400_000)