"""How much does each measured input move the answer?  Paired runs (same seeds) of model4 with one input changed.
-> leverage4.json.   Usage: python3 leverage4.py [n_per_run] [n_seeds]"""
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__))
VARIANTS = {
    "base": {},
    # this project's optical scaling-relation limits removed (far-infrared band unconstrained)
    "no optical galaxy limits (fp_kiii, tf_kiii)": dict(mid_curve=((0.75, 1e-12), (0.9, 1e-12))),
    "no galaxy surveys at all": dict(use_surveys=False),
    "no radio searches (seti_meta)": dict(use_seti=False),
    "loud beacons common": dict(f_beacon=(0.03, 1.0)),
    # planetary lifetime (mdwarf_early, planet_lifetime)
    "planets certainly die young (p_cap_short = 1)": dict(p_cap_short=1.0),
    "planets live as long as their stars (p_cap_short = 0)": dict(p_cap_short=0.0),
    "red-dwarf planets dead (p_M_class = 0)": dict(p_M_class=0.0),
    "red-dwarf planets alive (p_M_class = 1)": dict(p_M_class=1.0),
    # timing of life on Earth (timing, impact_window)
    "life dated from LUCA (tau/T 0.03-0.06)": dict(tau_over_T=(0.03, 0.06)),
    "life dated from accepted fossils (tau/T 0.16-0.20)": dict(tau_over_T=(0.16, 0.20)),
    # in-galaxy speed (flyby: short hops make slow settlement possible) -- fast only
    "settlement fronts at least 1% of light speed": dict(v_gal=(1e-2, 0.3)),
    "settlement fronts at most 0.1% of light speed": dict(v_gal=(1e-4, 1e-3)),
    # convergence / hard steps
    "no hard steps at all": dict(m_weights={0: 1.0}),
    "three or four hard steps": dict(m_weights={3: 0.5, 4: 0.5}),
}


def job(a):
    name, seed, n = a
    r = M.run(M.vary(M.BASE, n=n, seed=seed, **VARIANTS[name]))
    out = {}
    for s in ("site", "ssa_all", "ssa_obs"):
        q = r["schools"][s]
        out[s] = dict(alone=q["p_no_other_ever_in_galaxy"], stuck=q["posterior"]["stuck"], unnoticed=q["posterior"]["unnoticed"] + q["posterior"]["hidden"], ess=q["ess"])
    out["prior_alone"] = r["prior_stats"]["p_no_other_ever_in_galaxy"]
    return name, seed, out


if __name__ == "__main__":
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 200_000
    k = int(sys.argv[2]) if len(sys.argv) > 2 else 3
    jobs = [(name, 500 + i, n) for name in VARIANTS for i in range(k)]
    with Pool(int(os.environ.get("FERMI_PROCS", "4"))) as p:
        res = p.map(job, jobs, chunksize=1)
    by = {}
    for name, seed, out in res:
        by.setdefault(name, {})[seed] = out
    summary = {}
    for name in VARIANTS:
        row = {}
        for s in ("site", "ssa_all", "ssa_obs"):
            for f in ("alone", "stuck", "unnoticed"):
                v = np.array([by[name][sd][s][f] for sd in sorted(by[name])])
                b = np.array([by["base"][sd][s][f] for sd in sorted(by["base"])])
                row[f"{s}_{f}"] = float(v.mean()); row[f"{s}_{f}_delta"] = float((v - b).mean()); row[f"{s}_{f}_delta_sd"] = float((v - b).std(ddof=1) / np.sqrt(len(v))) if len(v) > 1 else None
        row["prior_alone"] = float(np.mean([by[name][sd]["prior_alone"] for sd in by[name]]))
        summary[name] = row
    json.dump(dict(n_per_run=n, seeds=k, variants={k_: {kk: (list(vv) if isinstance(vv, tuple) else vv) for kk, vv in v.items()} for k_, v in VARIANTS.items()}, summary=summary), open(os.path.join(HERE, "leverage4.json"), "w"), indent=1, default=str)
    for name, row in summary.items():
        print(f"{name:55s} alone {row['site_alone']:.3f} ({row['site_alone_delta']:+.3f})  stuck {row['site_stuck']:.3f} ({row['site_stuck_delta']:+.3f})  typical-all alone {row['ssa_all_alone']:.3f} ({row['ssa_all_alone_delta']:+.3f})")