#!/usr/bin/env python
"""Descriptive analysis of counterfactual labels: gain distribution, harm rate, variance decomposition of the
exhaustive 24-action audit into state and action components, per-action mean gain, oracle vs random vs STOP,
and the dependence on the prefix depth. Usage: python -m evaluation.label_analysis <labels_dir> [--out results/analyzed/labels_<name>.json]
"""
from __future__ import annotations

import argparse, json, sys
from pathlib import Path

import numpy as np
import torch


def analyse(path: Path) -> dict:
    d = torch.load(path, weights_only=False)
    recs = d["records"]
    g_all, r_pre_all = [], []
    aud = []
    for r in recs:
        v = r["valid"]
        g_all.append(r["gain"][v].numpy().ravel()); r_pre_all.append(np.full(int(v.sum()) * 4, r["r_pre"]))
        if "all_gains" in r:
            aud.append((r["all_gains"][v].numpy(), r["base_loss"][v].numpy(), r["r_pre"]))
    g = np.concatenate(g_all); rp = np.concatenate(r_pre_all)
    out = {"n_labels": int(g.size), "mean": float(g.mean()), "std": float(g.std()), "median": float(np.median(g)),
           "frac_pos": float((g > 0).mean()), "frac_neg": float((g < 0).mean()), "frac_abs_gt_0.01": float((np.abs(g) > 0.01).mean()),
           "frac_abs_gt_0.1": float((np.abs(g) > 0.1).mean()), "q05": float(np.quantile(g, 0.05)), "q95": float(np.quantile(g, 0.95)),
           "by_prefix_depth": {int(k): {"mean": float(g[rp == k].mean()), "frac_pos": float((g[rp == k] > 0).mean()), "n": int((rp == k).sum())} for k in np.unique(rp)}}
    if aud:
        G = np.concatenate([a for a, _, _ in aud])           # [n_states, A]
        base = np.concatenate([b for _, b, _ in aud])
        A = G.shape[1]
        state_mean = G.mean(1, keepdims=True)
        total_var = G.var()
        var_state = state_mean.var()                            # between-state variance of the mean gain
        var_action_within = (G - state_mean).var()              # within-state (action-specific) variance
        best = G.max(1); best_clip = np.clip(best, 0, None)
        rand = G.mean(1)
        out["audit"] = {
            "n_states": int(G.shape[0]), "n_actions": int(A),
            "var_total": float(total_var), "var_between_states": float(var_state), "var_within_state_actions": float(var_action_within),
            "within_share": float(var_action_within / max(total_var, 1e-12)),
            "oracle_one_step_mean": float(best_clip.mean()), "random_action_mean": float(rand.mean()),
            "stop_mean": 0.0, "frac_states_any_positive": float((best > 0).mean()),
            "frac_states_all_positive": float((G > 0).all(1).mean()), "frac_states_all_negative": float((G < 0).all(1).mean()),
            "mean_rank_corr_between_actions": None,
            "per_action_mean": G.mean(0).tolist(), "per_action_frac_best": np.bincount(G.argmax(1), minlength=A).tolist(),
            "base_loss_mean": float(base.mean()),
            "oracle_gain_relative_to_base": float(best_clip.mean() / max(base.mean(), 1e-9)),
            "single_vs_pair_mean": {"single": float(G[:, :16].mean()), "pair": float(G[:, 16:].mean())},
            "winners_curse_check_split_half": None,
        }
        # winner's-curse check: pick the best action on half the... not possible per-state (one label per action);
        # instead measure how often the best action of state i is also the best of a *different* state with the
        # same prefix depth (action consistency) vs chance 1/A.
        am = G.argmax(1)
        consist = (am[:-1] == am[1:]).mean()
        out["audit"]["consecutive_state_best_action_agreement"] = float(consist)
        # decomposition of the one-step oracle: continue/stop on the state mean vs choosing the best action
        state_level = np.clip(G.mean(1), 0, None)
        out["audit"]["oracle_state_level_mean"] = float(state_level.mean())       # STOP unless the average action helps
        out["audit"]["oracle_action_selection_increment"] = float(best_clip.mean() - state_level.mean())
        out["audit"]["random_with_perfect_stop_mean"] = float(state_level.mean())
        out["audit"]["chance_agreement"] = 1.0 / A
        # relation between base loss and oracle gain
        hi = base > np.median(base)
        out["audit"]["oracle_gain_high_vs_low_base_loss"] = {"high": float(best_clip[hi].mean()), "low": float(best_clip[~hi].mean())}
        # where the opportunity is, is a random action still good? (mean over the 24 actions vs their max)
        out["audit"]["random_gain_high_vs_low_base_loss"] = {"high": float(rand[hi].mean()), "low": float(rand[~hi].mean())}
        out["audit"]["frac_actions_positive_high_vs_low"] = {"high": float((G[hi] > 0).mean()), "low": float((G[~hi] > 0).mean())}
        out["audit"]["within_state_std_high_vs_low"] = {"high": float((G[hi] - G[hi].mean(1, keepdims=True)).std()), "low": float((G[~hi] - G[~hi].mean(1, keepdims=True)).std())}
        # top decile of base loss
        q = base >= np.quantile(base, 0.9)
        out["audit"]["top_decile_base_loss"] = {"oracle": float(best_clip[q].mean()), "random": float(rand[q].mean()), "frac_positive": float((G[q] > 0).mean()), "n": int(q.sum())}
    return out


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("labels_dir"); ap.add_argument("--out", default=None)
    args = ap.parse_args(argv)
    res = {}
    for part in ["train", "calib"]:
        p = Path(args.labels_dir) / f"labels_{part}.pt"
        if p.exists():
            res[part] = analyse(p)
            print(part, json.dumps({k: v for k, v in res[part].items() if k != "audit"}, indent=None)[:600])
            if "audit" in res[part]:
                a = res[part]["audit"]
                print("  audit:", {k: (round(v, 5) if isinstance(v, float) else v) for k, v in a.items() if k not in ("per_action_mean", "per_action_frac_best")})
    out = Path(args.out or f"results/analyzed/labels_{Path(args.labels_dir).name}.json")
    out.parent.mkdir(parents=True, exist_ok=True); out.write_text(json.dumps(res, indent=1))
    print("wrote", out)


if __name__ == "__main__":
    main()