#!/usr/bin/env python3
"""Rounds the registered outputs (results/R1.json, results/R2.json) for display, and adds the ranges
across records the paper quotes. Writes results/S.json; computes nothing new."""
import json, sys
from pathlib import Path

ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else ".")
R1 = json.loads((ROOT / "results/R1.json").read_text())
R2 = json.loads((ROOT / "results/R2.json").read_text())
P = json.loads((ROOT / "data/parameters.json").read_text())
NAMES = P["datasets"]


def sig(x, n=2):
    return float(f"{x:.{n}g}")


def dp(x, n):
    return round(float(x), n)


S = {"null_sims": P["null_sims"], "observed_sims": P["observed_sims"], "qtable_sims": P["qtable_sims"],
     "test_start": P["test_years"][0], "test_end": P["test_years"][1], "records": len(NAMES),
     "max_abs_dev_from_published_c": sig(max(R1[d]["max_abs_dev_from_published_adjusted"] for d in NAMES), 1),
     "levels": {"alpha": 0.05, "strict": 0.01, "quadratic": 0.02},
     "datasets": {}, "ranges": {}, "decision": R2["decision"]}
for d in NAMES:
    a, u = R1[d]["adjusted"], R1[d]["unadjusted"]
    row = {
        "tmax": dp(a["tmax"], 2), "change_year": a["change_year"], "phi": dp(a["phi"], 2),
        "p_nominal": sig(a["p_tmax"]), "q95_nominal": dp(a["q95"], 2), "q99_nominal": dp(a["q99"], 2),
        "zq_monthly": dp(a["zq_monthly"], 2), "zq_annual": dp(a["zq_annual"], 2),
        "p_quad_monthly_nominal": sig(a["p_quad_monthly"]), "p_quad_annual_nominal": sig(a["p_quad_annual"]),
        "raw_tmax": dp(u["tmax"], 2), "raw_p": sig(u["p_tmax"]),
    }
    for m in P["noise_models"]:
        x = R2["noise_models"][m][d]
        row[m] = {
            "size05": dp(x["size_tmax_005"], 3), "size01": dp(x["size_tmax_001"], 3),
            "q95": dp(x["tmax_q95"], 2), "q99": dp(x["tmax_q99"], 2), "p_tmax": sig(x["p_tmax"]),
            "size_quad_monthly05": dp(x["size_quad_monthly_005"], 3), "size_quad_annual05": dp(x["size_quad_annual_005"], 3),
            "p_quad_monthly": sig(x["p_quad_monthly"]), "p_quad_annual": sig(x["p_quad_annual"]),
            **({"d": dp(x["noise"]["d"], 2)} if "d" in x["noise"] else {"ar": dp(x["noise"]["ar"], 2), "ma": dp(x["noise"]["ma"], 2)}),
        }
    S["datasets"][d] = row

rows = S["datasets"]


def rng(key, model=None):
    v = [rows[d][model][key] if model else rows[d][key] for d in NAMES]
    return min(v), max(v)


ranges = {}
for key in ("tmax", "p_nominal", "zq_monthly", "zq_annual", "raw_p"):
    ranges[key + "_min"], ranges[key + "_max"] = rng(key)
for m in P["noise_models"]:
    for key in ("size05", "size01", "q95", "q99", "p_tmax", "size_quad_monthly05", "size_quad_annual05", "p_quad_monthly", "p_quad_annual"):
        ranges[f"{m}_{key}_min"], ranges[f"{m}_{key}_max"] = rng(key, m)
    ranges[f"{m}_records_p_tmax_below_001"] = sum(R2["noise_models"][m][d]["p_tmax"] < 0.01 for d in NAMES)
    ranges[f"{m}_records_p_tmax_below_005"] = sum(R2["noise_models"][m][d]["p_tmax"] < 0.05 for d in NAMES)
    ranges[f"{m}_records_p_quad_monthly_below_002"] = sum(R2["noise_models"][m][d]["p_quad_monthly"] < 0.02 for d in NAMES)
    ranges[f"{m}_records_p_quad_annual_below_002"] = sum(R2["noise_models"][m][d]["p_quad_annual"] < 0.02 for d in NAMES)
    ranges[f"{m}_records_p_quad_monthly_below_005"] = sum(R2["noise_models"][m][d]["p_quad_monthly"] < 0.05 for d in NAMES)
ranges["long_memory_d_min"], ranges["long_memory_d_max"] = rng("d", "long_memory")
ranges["records_raw_p_below_005"] = sum(R1[d]["unadjusted"]["p_tmax"] < 0.05 for d in NAMES)
ranges["change_year_min"], ranges["change_year_max"] = rng("change_year")
S["ranges"] = ranges
(ROOT / "results/S.json").write_text(json.dumps(S, indent=1, sort_keys=True) + "\n")
print(json.dumps(S["ranges"], indent=1))
