"""Ledger, settling, common window and Morison fits for Dütsch cases (numpy only).

    python analyze.py CASE_DIR [CASE_DIR ...]   # cases of one comparison
"""
from __future__ import annotations

import json
import math
import os
import sys

import numpy as np

import protocol as p


def load(path):
    with open(os.path.join(path, "case.json")) as fh:
        case = json.load(fh)
    raw = np.load(os.path.join(path, "raw.npz"))
    case.update(force=raw["force"], diag=raw["diag"], samples=raw["samples"], path=path)
    return case


def cycle_slices(case):
    """{k: slice of step indices m whose interval [m, m+1] lies in paper cycle k}."""
    T, q = case["T"], case["T"] // 4
    return {k: slice(q + (k - 1) * T, q + k * T)
            for k in range(1, case["cycles"] + 1) if q + k * T <= len(case["force"])}


def coarse_wave(case, k, r=1):
    """Normalized per-step Fx of cycle k, averaged over r native steps per sample."""
    f = case["force"][cycle_slices(case)[k], 0] / case["F_ref"]
    return f.reshape(-1, r).mean(axis=1)


def profiles(case, k):
    """{(phase, station, x2): (u, v)} of cycle k, time-averaged over the bin steps
    (equal weights per step) from valid samples only."""
    s = case["samples"]
    acc = {}
    for _, _, ph, xs, x2, u, v, _ in s[s[:, 1] == k]:
        if not (math.isnan(u) or math.isnan(v)):
            acc.setdefault((ph, xs, x2), []).append((u, v))
    return {key: tuple(np.mean(vals, axis=0)) for key, vals in acc.items()}


def coverage(case, k):
    s = case["samples"]
    rows = s[s[:, 1] == k]
    missing = int(np.isnan(rows[:, 5]).sum())
    return dict(samples=len(rows), missing=missing, bin_steps=len(set(rows[:, 0])),
                min_fluid_weight=float(np.nanmin(np.where(rows[:, 7] > 0, rows[:, 7], np.nan)))
                if len(rows) else math.nan)


def common_points(dicts):
    keys = set(dicts[0])
    for d in dicts[1:]:
        keys &= set(d)
    return sorted(keys)


def candidates(case):
    """Per-step Fx of each reported force candidate (none is selected as an RHS)."""
    d = case["diag"]
    return {"cme": d[:, 0], "cme_wen": d[:, 0] + d[:, 2], "cme_node": d[:, 0] + d[:, 6] - d[:, 9]}


def cycle_morison(case, ks, fx=None):
    T = case["T"]
    fx = case["force"][:, 0] if fx is None else fx
    idx = np.concatenate([np.arange(len(fx))[cycle_slices(case)[k]] for k in ks])
    b = np.array([p.interval_bases(int(m), T) for m in idx])
    return p.morison_fit(fx[idx] / case["F_ref"], b[:, 0], b[:, 1])


def pair_metrics(case, k):
    a, b = coarse_wave(case, k - 1), coarse_wave(case, k)
    pa, pb = profiles(case, k - 1), profiles(case, k)
    common = common_points([pa, pb])  # empty set -> inf -> pair fails (FAIL-CLOSED)
    prof = max((max(abs(pa[c][0] - pb[c][0]), abs(pa[c][1] - pb[c][1])) for c in common),
               default=math.inf) / case["u_max"]
    (cd0, ci0), (cd1, ci1) = cycle_morison(case, [k - 1]), cycle_morison(case, [k])
    m = dict(wave=float(np.sqrt(np.mean((b - a) ** 2))),
             scalar=float(max(abs(b.mean() - a.mean()), abs(np.sqrt(np.mean(b * b)) - np.sqrt(np.mean(a * a))),
                              abs(np.abs(b).max() - np.abs(a).max()))),
             profile=float(prof), morison=float(max(abs(cd1 - cd0), abs(ci1 - ci0))),
             profile_points=len(common))
    m["pass"] = all(m[key] <= tol for key, tol in p.SETTLE.items())
    return m


def ledger(case):
    ks = sorted(cycle_slices(case))
    pairs = {k: pair_metrics(case, k) for k in ks if k - 1 > p.DISCARD}
    return pairs, p.settled_start({k: v["pass"] for k, v in pairs.items()})


def main(argv):
    """argv: CASE_DIR ... [--reference CASE_DIR]; the reference defaults to the
    largest D (resolution ladder); domain runs name the 60 D case explicitly."""
    ref_path = None
    if "--reference" in argv:
        i = argv.index("--reference")
        ref_path, argv = argv[i + 1], argv[:i] + argv[i + 2:]
    paths = argv if ref_path is None or ref_path in argv else argv + [ref_path]
    cases = [load(path) for path in paths]
    ledgers = [ledger(c) for c in cases]
    window = p.common_window([k for _, k in ledgers],
                             [{k: v["pass"] for k, v in pairs.items()} for pairs, _ in ledgers])
    report = dict(window=window, cases=[])
    for case, (pairs, k_case) in zip(cases, ledgers):
        entry = {key: case[key] for key in ("path", "D", "T", "box", "n", "u_max", "nu", "tau", "A",
                                            "steps", "cycles", "F_ref")}
        entry.update(duration_T=case["steps"] / case["T"], K_case=k_case, pairs=pairs,
                     coverage={k: coverage(case, k) for k in sorted(cycle_slices(case)) if k > p.DISCARD},
                     no_donor=float(case["diag"][:, 13].sum()), eq_only=float(case["diag"][:, 12].sum()))
        if window:  # residual drift inside the scored window (pairs K..K+2)
            entry["window_drift"] = {key: max(pairs[k][key] for k in window)
                                     for key in ("wave", "profile", "morison")}
            entry["morison_obs"] = cycle_morison(case, window)
            entry["morison"] = {name: cycle_morison(case, window, fx)
                                for name, fx in candidates(case).items()}
            entry["fy_mean"] = float(np.mean(np.concatenate(
                [case["force"][cycle_slices(case)[k], 1] for k in window])) / case["F_ref"])
        report["cases"].append(entry)
    if window and len(cases) > 1:  # cross-resolution waveform against the finest case
        fine = (next(c for c in cases if c["path"] == ref_path) if ref_path
                else max(cases, key=lambda c: c["D"]))
        report["reference"] = fine["path"]
        ref = np.concatenate([coarse_wave(fine, k) for k in window])
        for case, entry in zip(cases, report["cases"]):
            r = fine["T"] // case["T"]  # 1 for equal-D domain runs
            ref_c = ref.reshape(-1, r).mean(axis=1)
            own = np.concatenate([coarse_wave(case, k) for k in window])
            entry["wave_rms_vs_finest"] = float(np.sqrt(np.mean((own - ref_c) ** 2)))
        for k in window:  # profiles on the common valid point set of all compared cases
            per_case = [profiles(c, k) for c in cases]
            keys = common_points(per_case)
            report.setdefault("common_profile_points", {})[k] = len(keys)
            if not keys:
                report["verdict"] = "FAIL-CLOSED: empty common profile set"
                continue
            fine_p = per_case[cases.index(fine)]
            for case, entry, prof in zip(cases, report["cases"], per_case):
                d = max(max(abs(prof[key][0] - fine_p[key][0]), abs(prof[key][1] - fine_p[key][1]))
                        for key in keys) / case["u_max"]
                entry.setdefault("profile_max_vs_finest", {})[k] = float(d)
        ref_entry = report["cases"][cases.index(fine)]
        for entry in report["cases"]:  # conservative: difference + both drifts
            if entry is ref_entry:
                continue
            drift = {key: entry["window_drift"][key] + ref_entry["window_drift"][key]
                     for key in ("wave", "profile", "morison")}
            dm = max(abs(entry["morison"]["cme"][0] - ref_entry["morison"]["cme"][0]),
                     abs(entry["morison"]["cme"][1] - ref_entry["morison"]["cme"][1]))
            entry["conservative"] = dict(
                wave=entry["wave_rms_vs_finest"] + drift["wave"],
                profile=max(entry.get("profile_max_vs_finest", {}).values(), default=math.inf)
                + drift["profile"],
                morison_cme=dm + drift["morison"])
            bands = p.DOMAIN if ref_path else p.CROSS  # domain runs name the 60 D reference
            c = entry["conservative"]
            entry["pass_scope"] = ("conservative (|diff| + both window drifts) on the obs/CME Fx "
                                   "waveform, CME Morison fit and simulation-to-simulation profiles "
                                   "only; CME+Wen and CME+node are reported, not checked; no "
                                   "Fig. 8 / literature verdict")
            entry["pass"] = (c["wave"] <= bands["wave"] and c["profile"] <= bands["profile"]
                             and ("morison" not in bands or c["morison_cme"] <= bands["morison"]))
    print(json.dumps(report, indent=1, default=float))


if __name__ == "__main__":
    main(sys.argv[1:])