"""environment320-v1 pure all-assigned aggregation (no I/O; import has no side effects).

API: aggregate(results, assigned_units). assigned_units is the complete roster from create_binding().units()
(320 episodes + 4 probes, validated); results are binding.score() dicts (list with optional None entries, or a
dict unit_id -> dict|None). Assigned episodes with no result become NOT_RUN (known False, common unknown bounds);
unreached probes are NOT_RUN diagnostics, never episodes. Unknown or duplicate IDs are refused (RosterError).
Every assigned episode stays in its arm's denominator (160 per arm). Unknown (technical, input-bound or NOT_RUN)
episodes enter as executable-domain bounds, never as imputed values. Known-only means are
reported alongside, labelled as such. Normalization uses the design constant 520000 cents, never an observed
maximum. No significance test, stopping rule or continuation decision depends on these numbers.
"""
import importlib.util
import math
from decimal import Decimal
from pathlib import Path

_spec = importlib.util.spec_from_file_location("_environment320_env_core_analysis", Path(__file__).resolve().parent / "env_core.py")
core = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(core)

METRICS = ("L", "P", "no_deal", "A", "emissions_kg", "buyer_gain_cents", "seller_gain_cents", "joint_cents")
ARMS = ("disclosed", "undisclosed")


def _mean(xs):
    return sum(xs) / len(xs) if xs else None


def _sd(xs):
    if len(xs) < 2:
        return None
    m = _mean(xs)
    return math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1))


def _bounds():
    return core.unknown_bounds()


def arm_summary(rows):
    n = len(rows)
    known = [r for r in rows if r["known"]]
    u = n - len(known)
    b = _bounds()
    out = {"assigned": n, "known": len(known), "unknown": u,
           "status_counts": {s: sum(1 for r in rows if r["status"] == s) for s in sorted({r["status"] for r in rows})}}
    for m in METRICS:
        vals = [r["outcome"][m] for r in known]
        lo, hi = b[m]
        s = sum(vals)
        out[m] = {"known_mean": _mean(vals), "known_sd": _sd(vals),
                  "all_assigned_bounds": [(s + lo * u) / n, (s + hi * u) / n] if n else None}
    for r in known:
        o = r["outcome"]
        if o["L"] != o["P"] + core.NO_DEAL_L * o["no_deal"]:
            raise AssertionError(f"decomposition:{r['unit_id']}")
    out["L_norm_all_assigned_bounds"] = [x / core.L_RANGE for x in out["L"]["all_assigned_bounds"]] if n else None
    deals = [r for r in known if r["status"] == "DEAL"]
    prices = [r["outcome"]["price_cents"] for r in deals]
    out["prices"] = {"deals": len(deals), "mean_cents": _mean(prices), "sd_cents": _sd(prices),
                     "min_cents": min(prices) if prices else None, "max_cents": max(prices) if prices else None}
    out["plan_counts"] = {pid: sum(1 for r in deals if r["outcome"]["plan"] == pid) for pid in core.EMISSIONS_KG}
    out["lower_tail"] = {}
    for role in ("buyer", "seller"):
        g = sorted(r["outcome"][f"{role}_gain_cents"] for r in known)
        out["lower_tail"][role] = {
            "known_below_fallback": sum(1 for x in g if x < 0), "known_min_gain_cents": g[0] if g else None,
            "known_p10_gain_cents_nearest_rank": g[max(0, math.ceil(0.1 * len(g)) - 1)] if g else None,
            "unknown_could_be_below_fallback": u}
    out["codes"] = {k: sum(1 for r in known if r["codes"][k]) for k in
                    ("accepter_below_fallback", "proposer_below_fallback", "dominated_plan", "non_DO_package_joint_loss")}
    return out


def cell_spread(rows):
    cells = {}
    for r in rows:
        cells.setdefault(r["cell"], []).append(r)
    out = {}
    for c in sorted(cells):
        known = [r for r in cells[c] if r["known"]]
        out[c] = {"assigned": len(cells[c]), "known": len(known)}
        for m in ("L", "A", "buyer_gain_cents", "seller_gain_cents"):
            vals = [r["outcome"][m] for r in known]
            out[c][m] = {"mean": _mean(vals), "sd": _sd(vals), "min": min(vals) if vals else None,
                         "max": max(vals) if vals else None}
    return out


class RosterError(ValueError):
    """Unidentifiable roster or result data: refused, never guessed."""


ROSTER_EPISODES, ROSTER_PROBES, PER_ARM, PER_CELL = 320, 4, 160, 40


def validate_roster(assigned_units):
    """The complete assigned roster (binding.create_binding().units(), passed by the caller): exactly 320 episode
    and 4 probe units, unique IDs, 160 episodes per arm, 40 per exact cell, 4 distinct arm x role probes."""
    if not isinstance(assigned_units, list) or not all(isinstance(u, dict) for u in assigned_units):
        raise RosterError("roster_not_list_of_units")
    ids = [u.get("unit_id") for u in assigned_units]
    if any(not isinstance(i, str) for i in ids) or len(set(ids)) != len(ids):
        raise RosterError("roster_ids_missing_or_duplicate")
    eps = [u for u in assigned_units if u.get("kind") == "episode"]
    prs = [u for u in assigned_units if u.get("kind") == "probe"]
    if len(eps) != ROSTER_EPISODES or len(prs) != ROSTER_PROBES or len(eps) + len(prs) != len(assigned_units):
        raise RosterError("roster_size")
    for u in eps:
        d = u.get("design") or {}
        if d.get("arm") not in ARMS or d.get("cell") != f"{d.get('arm')}|{d.get('opener')}|{d.get('row_order')}":
            raise RosterError(f"roster_episode_design:{u['unit_id']}")
    if any(sum(1 for u in eps if u["design"]["arm"] == a) != PER_ARM for a in ARMS):
        raise RosterError("roster_arm_counts")
    cells = {}
    for u in eps:
        cells[u["design"]["cell"]] = cells.get(u["design"]["cell"], 0) + 1
    if len(cells) != 8 or any(v != PER_CELL for v in cells.values()):
        raise RosterError("roster_cell_counts")
    if len({(u.get("design", {}).get("arm"), u.get("design", {}).get("role")) for u in prs}) != ROSTER_PROBES:
        raise RosterError("roster_probes")
    return {u["unit_id"]: u for u in assigned_units}


def not_run_episode(u):
    """An assigned episode with no scoreable result (never admitted, or incomplete with no service terminal):
    status NOT_RUN, known False, the same executable-domain unknown bounds as any unknown episode; never imputed."""
    d = u["design"]
    return {"study": "environment320", "protocol_version": core.PROTOCOL_VERSION, "unit_id": u["unit_id"],
            "kind": "episode", "arm": d["arm"], "opener": d["opener"], "row_order": d["row_order"], "cell": d["cell"],
            "ordinal": d.get("ordinal"), "block": d.get("block"), "status": "NOT_RUN", "known": False,
            "outcome": None, "codes": None, "bounds": {k: list(v) for k, v in core.unknown_bounds().items()},
            "checks": {r: {"role": r, "status": "CHECK_NOT_RUN", "reason": None, "answer": None, "correct": None}
                       for r in ("buyer", "seller")}}


def not_run_probe(u):
    d = u["design"]
    return {"unit_id": u["unit_id"], "kind": "probe", "arm": d["arm"], "role": d["role"], "probe": None,
            "status": "NOT_RUN"}


def complete_results(results, assigned_units):
    """Map caller results onto the full roster. results: list of binding.score() dicts (None entries allowed and
    ignored, since the roster identifies absent units) or a dict unit_id -> dict|None. Unknown IDs, duplicates,
    kind/arm/cell mismatches and malformed known outcomes are refused."""
    roster = validate_roster(assigned_units)
    if isinstance(results, dict):
        items = []
        for uid, r in results.items():
            if r is not None and r.get("unit_id") != uid:
                raise RosterError(f"result_key_mismatch:{uid}")
            if uid not in roster:
                raise RosterError(f"unknown_unit_id:{uid}")
            items.append(r)
    elif isinstance(results, list):
        items = results
    else:
        raise RosterError("results_not_list_or_dict")
    got = {}
    for r in items:
        if r is None:
            continue
        if not isinstance(r, dict) or not isinstance(r.get("unit_id"), str):
            raise RosterError("result_without_unit_id")
        uid = r["unit_id"]
        if uid not in roster:
            raise RosterError(f"unknown_unit_id:{uid}")
        if uid in got:
            raise RosterError(f"duplicate_result:{uid}")
        u = roster[uid]
        if r.get("kind") != u["kind"] or r.get("arm") != u["design"]["arm"]:
            raise RosterError(f"result_design_mismatch:{uid}")
        if u["kind"] == "episode":
            if r.get("cell") != u["design"]["cell"] or not isinstance(r.get("known"), bool):
                raise RosterError(f"result_design_mismatch:{uid}")
            if r["known"] and not isinstance(r.get("outcome"), dict):
                raise RosterError(f"known_result_without_outcome:{uid}")
            if not r["known"] and r.get("outcome") is not None:
                raise RosterError(f"unknown_result_with_outcome:{uid}")
        got[uid] = r
    out = []
    for uid, u in roster.items():
        if uid in got:
            out.append(got[uid])
        else:
            out.append(not_run_episode(u) if u["kind"] == "episode" else not_run_probe(u))
    return out


def aggregate(results, assigned_units):
    """All-assigned aggregation over the complete roster (see complete_results). The caller must pass
    assigned_units = create_binding().units(); there is no default and no import of the binding here."""
    full = complete_results(results, assigned_units)
    episodes = [r for r in full if r["kind"] == "episode"]
    probes = [r for r in full if r["kind"] == "probe"]
    by = {a: [r for r in episodes if r["arm"] == a] for a in ARMS}
    arms = {a: arm_summary(by[a]) for a in ARMS}
    T, C = arms["disclosed"], arms["undisclosed"]
    diff = {}
    for m in METRICS:
        tb, cb = T[m]["all_assigned_bounds"], C[m]["all_assigned_bounds"]
        tk, ck = T[m]["known_mean"], C[m]["known_mean"]
        se = None
        if T[m]["known_sd"] is not None and C[m]["known_sd"] is not None and T["known"] and C["known"]:
            se = math.sqrt(T[m]["known_sd"] ** 2 / T["known"] + C[m]["known_sd"] ** 2 / C["known"])
        diff[m] = {"all_assigned_bounds": [tb[0] - cb[1], tb[1] - cb[0]] if tb and cb else None,
                   "known_difference": (tk - ck) if tk is not None and ck is not None else None,
                   "known_se_descriptive": se}
    # exact decomposition of the all-assigned difference when every episode is known
    if T["unknown"] == 0 and C["unknown"] == 0:
        dl, dp, dn = (diff[m]["all_assigned_bounds"][0] for m in ("L", "P", "no_deal"))
        diff["decomposition_check"] = {"dL": dl, "dP_plus_520000_dNoDeal": dp + core.NO_DEAL_L * dn,
                                       "holds": abs(dl - (dp + core.NO_DEAL_L * dn)) < 1e-6}
    bal = {m: balanced_difference(by["disclosed"], by["undisclosed"], m) for m in ("L", "A")}
    diff["cell_balanced"] = bal
    diff["usd_per_tonne_descriptive"] = ratio({"dL_balanced_cents": bal["L"]["difference"],
                                               "dA_balanced_kg": bal["A"]["difference"],
                                               "dA_balanced_se": bal["A"]["se"],
                                               "dA_all_assigned_bounds": diff["A"]["all_assigned_bounds"]})
    supplied = len([r for r in (results.values() if isinstance(results, dict) else results) if r is not None])
    return {"protocol_version": core.PROTOCOL_VERSION, "arms": arms, "difference_disclosed_minus_undisclosed": diff,
            "roster": {"episodes_assigned": len(episodes), "probes_assigned": len(probes), "results_supplied": supplied,
                       "episodes_not_run": sum(1 for r in episodes if r["status"] == "NOT_RUN"),
                       "probes_not_run": sum(1 for p in probes if p.get("status") == "NOT_RUN")},
            "cell_spread": cell_spread(episodes),
            "probes": [{"unit_id": p["unit_id"], "arm": p["arm"], "role": p["role"],
                        "status": p["probe"]["status"] if p["probe"] else "NOT_RUN",
                        "correct": p["probe"]["correct"] if p["probe"] else None} for p in probes],
            "checks": check_summary(episodes),
            "notes": ["Probes and checks are descriptive and never determine outcome-based continuation.",
                      "Ratio is 1000 * dL_usd / dA_kg, descriptive only; no WTP parameter is identified."]}


def balanced_arm(rows, metric):
    """Prespecified scenario-cell-balanced known-only arm mean: equal-weight average of the arm's 4 exact-cell
    means (opener x row order); SE = sqrt(sum_c sd_c^2 / n_c) / 4. None if any cell has < 2 known episodes."""
    cells = {}
    for r in rows:
        if r["known"]:
            cells.setdefault(r["cell"], []).append(r["outcome"][metric])
    if len(cells) != 4 or any(len(v) < 2 for v in cells.values()):
        return {"mean": None, "se": None, "cells": len(cells)}
    means = [_mean(v) for v in cells.values()]
    var = sum(_sd(v) ** 2 / len(v) for v in cells.values())
    return {"mean": sum(means) / 4, "se": math.sqrt(var) / 4, "cells": 4}


def balanced_difference(t_rows, c_rows, metric):
    t, c = balanced_arm(t_rows, metric), balanced_arm(c_rows, metric)
    if t["mean"] is None or c["mean"] is None:
        return {"difference": None, "se": None, "disclosed": t, "undisclosed": c}
    return {"difference": t["mean"] - c["mean"], "se": math.sqrt(t["se"] ** 2 + c["se"] ** 2), "disclosed": t, "undisclosed": c}


def ratio(comp):
    """Resolution C5: 1000 * dL(USD) / dA(kg), owner dollars per tonne averted, descriptive only.
    Withheld (value None, components and reason still reported) if the all-assigned missing-outcome dA interval is
    missing or includes zero, the balanced dA/SE is unavailable, or abs(dA) <= 2 * balanced SE.
    An interpretability guard only: never a stopping, continuation or significance rule."""
    out = {"value": None, "components": dict(comp)}
    ab = comp.get("dA_all_assigned_bounds")
    dl, da, se = comp.get("dL_balanced_cents"), comp.get("dA_balanced_kg"), comp.get("dA_balanced_se")
    if ab is None:
        return dict(out, reason="dA_all_assigned_bounds_missing")
    if ab[0] <= 0 <= ab[1]:
        return dict(out, reason="dA_all_assigned_bounds_include_zero")
    if dl is None or da is None or se is None:
        return dict(out, reason="balanced_components_unavailable")
    if abs(da) <= 2 * se:
        return dict(out, reason="abs_dA_not_greater_than_2_balanced_se")
    v = Decimal(repr(dl)) / 100 * 1000 / Decimal(repr(da))
    return dict(out, value=str(v.quantize(Decimal("0.01"))), reason="defined_descriptive_only")


def check_summary(episodes):
    out = {}
    for a in ARMS:
        for role in ("buyer", "seller"):
            cs = [r["checks"][role] for r in episodes if r["arm"] == a]
            st = {}
            for c in cs:
                st[c["status"]] = st.get(c["status"], 0) + 1
            answered = [c for c in cs if c["status"] == "CHECK_ANSWERED"]
            out[f"{a}|{role}"] = {"status_counts": st, "all_three_correct": sum(1 for c in answered if all(c["correct"].values())),
                                  "assigned": len(cs)}
    return out