"""Joint prior scan for model4: every disputed prior is redrawn at once, centred on BASE. -> integrate4_scan.json
Usage: python3 scan4.py [n_draws] [n_per_draw]"""
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__))
S3 = ("site", "ssa_all", "ssa_now", "ssa_obs")


def draw(k):
    """Round 3: the ranges now reach the optimists. The prior that life starts easily is uniform on 0.05-0.95 and the weight on
    'no candidate step at all' uniform on 0.1-0.8; the slow tail may be twice as long as mine; knobs held fixed before are varied."""
    rng = np.random.default_rng(5000 + k)
    lg = lambda lo, hi: float(10 ** rng.uniform(np.log10(lo), np.log10(hi)))
    p0 = rng.uniform(0.1, 0.8); ratio = rng.uniform(0.35, 1.1)
    tail = np.array([ratio ** j for j in range(1, 6)]); tail = (1 - p0) * tail / tail.sum()
    return dict(
        w_fast=float(rng.uniform(0.05, 0.95)), slow_decades=(3.0, lg(10, 200)), fast_decades=float(rng.uniform(2.0, 4.0)),
        m_weights={0: float(p0), **{j + 1: float(tail[j]) for j in range(5)}},
        step_decades=(2.0, float(rng.uniform(4, 12))), f_gate=(lg(0.003, 0.3), 1.0),
        q_barrier=lg(0.03, 0.4), W_exp=(float(rng.uniform(0.5, 2.0)), lg(4, 16)),
        q_cc=lg(0.01, 0.2), p_stealth=lg(0.003, 0.1), p_here=(1e-3, lg(0.03, 1.0)), eps_pocket=(1e-6, lg(0.01, 0.5)),
        c_claim=("uniform", 0.0, 1.0) if rng.random() < 0.8 else ("loguniform", lg(1e-6, 0.1), 1.0),
        p_harvest=float(rng.uniform(0.1, 0.9)), p_M_class=float(rng.uniform(0.05, 0.6)),
        p_cap_short=float(rng.uniform(0.0, 1.0)), q_ig=float(rng.uniform(0.2, 0.9)), p_persist_hazard=float(rng.uniform(0, 1)),
        f_site=(lg(1e-3, 0.1), 1.0), f_persist=(lg(3e-3, 0.3), 1.0), f_c=(lg(1e-4, 0.1), 1.0),
        tau_over_T=(0.02, float(rng.uniform(0.06, 0.20))), K_uv=(lg(0.01, 1.0), 1.0),
    )


def job(args):
    k, n = args
    kw = draw(k)
    r = M.run(M.vary(M.BASE, n=n, seed=9000 + k, **kw))
    row = dict(params={a: (list(b) if isinstance(b, tuple) else b) for a, b in kw.items()}, prior=r["prior"],
               prior_alone=r["prior_stats"]["p_no_other_ever_in_galaxy"])
    for s in S3:
        q = r["schools"][s]
        row[s] = dict(posterior=q["posterior"], ess=q["ess"], alone=q["p_no_other_ever_in_galaxy"], now=q["p_anyone_now_in_galaxy"], p_fast=q["p_fast"],
                      stuck_barrier=q["sub"]["stuck_barrier"], sterile=q["sub"]["alone_sterile"], p_barrier=q["p_barrier"])
    row["prior_fast"] = r["prior_stats"]["p_fast"]
    return row


def _job_k(a):
    return a[0], job(a)


if __name__ == "__main__":
    nd = int(sys.argv[1]) if len(sys.argv) > 1 else 400
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 100_000
    # resumable: finished draws are kept in a partial file, so an interrupted scan continues where it stopped
    part = os.path.join(HERE, f"integrate4_scan_partial_{n}.json")
    done = {r["k"]: r for r in json.load(open(part))} if os.path.exists(part) else {}
    todo = [(k, n) for k in range(nd) if k not in done]
    with Pool(int(os.environ.get("FERMI_PROCS", "4"))) as pool:
        for i, (k, row) in enumerate(pool.imap_unordered(_job_k, todo)):
            row["k"] = k; done[k] = row
            if i % 8 == 7 or i == len(todo) - 1:
                json.dump(list(done.values()), open(part, "w")); print("draws done", len(done), flush=True)
    rows = [done[k] for k in range(nd)]
    json.dump(rows, open(os.path.join(HERE, "integrate4_scan.json"), "w"))
    for s in S3:
        a = np.array([r[s]["alone"] for r in rows]); st = np.array([r[s]["posterior"]["stuck"] for r in rows])
        un = np.array([r[s]["posterior"]["unnoticed"] + r[s]["posterior"]["hidden"] for r in rows])
        f = lambda v: "/".join(f"{x:.2f}" for x in np.quantile(v, [0.05, 0.5, 0.95]))
        nb = np.array([r[s]["posterior"]["empty_universe"] + r[s]["posterior"]["rare"] for r in rows])
        print(f"{s:8s} alone {f(a)}  stuck {f(st)}  unnoticed+hidden {f(un)}  min ESS {min(r[s]['ess'] for r in rows):.0f}  "
              f"share of draws with alone > 1/2: {np.mean(a > 0.5):.2f}; with 'nobody' the leading answer: {np.mean((nb > st) & (nb > 1 - nb - st)):.2f}")
    pa = np.array([r["prior_alone"] for r in rows]); sa = np.array([r["site"]["alone"] for r in rows])
    print("prior alone 5/50/95:", np.quantile(pa, [0.05, 0.5, 0.95]).round(2), " posterior minus prior (site):", np.quantile(sa - pa, [0.05, 0.5, 0.95]).round(3),
          " share of draws where the observations raise P(alone):", float(np.mean(sa > pa)))