#!/usr/bin/env python3
"""Tier 2 behavioural exit check (Plan 5 Task 4).

Consumes the aggregated metrics JSON produced by
`scripts/aggregate-multi-seed-metrics.py` (Plan 5 Task 1B.2).

Pass criteria (per spec §2 Tier 2 — behavioural):
  1. `val_trades_per_bar` mean ≥ 0.005 on stable epochs (model trades).
  2. `val_active_frac` mean > 0.2 on stable epochs (model is not always Hold).
  3. Direction-argmax entropy > 0.8 × log(4) ≈ 1.109 on stable epochs.
     Preferred source: a `val_dir_entropy` aggregate metric. Fallback:
     derive from `val_dir_dist_*` (per-direction probability) if present.
     If neither is in the aggregate, FAIL with an explanatory message so
     the missing metric gets added to HEALTH_DIAG in a follow-up run.

Exits 0 on PASS, 1 on FAIL. Each check prints a `PASS:` / `FAIL:` line.
"""
from __future__ import annotations

import argparse
import json
import math
import sys
from statistics import mean


# 0.8 * log(4) ≈ 1.10903549
DIR_ENTROPY_FLOOR = 0.8 * math.log(4.0)


def _stable(series, warmup_end):
    return [x for x in series if x["epoch"] >= warmup_end]


def check_trades_per_bar(agg, warmup_end):
    series = agg.get("val_trades_per_bar")
    if not series:
        return False, (
            "val_trades_per_bar missing from aggregate — required for Tier 2"
        )
    stable = _stable(series, warmup_end)
    if not stable:
        return False, (
            f"No stable epochs for val_trades_per_bar past warmup_end={warmup_end}"
        )
    m = mean(x["mean"] for x in stable)
    if m < 0.005:
        return False, f"val_trades_per_bar mean = {m:.5f} < 0.005"
    return True, f"val_trades_per_bar mean = {m:.5f} >= 0.005"


def check_active_frac(agg, warmup_end):
    series = agg.get("val_active_frac")
    if not series:
        return False, (
            "val_active_frac missing from aggregate — required for Tier 2"
        )
    stable = _stable(series, warmup_end)
    if not stable:
        return False, (
            f"No stable epochs for val_active_frac past warmup_end={warmup_end}"
        )
    m = mean(x["mean"] for x in stable)
    if m <= 0.2:
        return False, f"val_active_frac mean = {m:.4f} <= 0.2"
    return True, f"val_active_frac mean = {m:.4f} > 0.2"


def _entropy_from_dist(probs):
    """Shannon entropy in nats; ignores non-positive probabilities."""
    return -sum(p * math.log(p) for p in probs if p > 0)


def check_dir_entropy(agg, warmup_end):
    """Direction argmax entropy > 0.8 × log(4).

    Prefers a pre-computed `val_dir_entropy` aggregate. Falls back to
    deriving entropy from per-direction probability mass (`val_dir_dist_short`
    / `val_dir_dist_hold` / `val_dir_dist_long` / `val_dir_dist_flat` —
    naming kept loose to match whatever HEALTH_DIAG ends up emitting).
    """
    direct = agg.get("val_dir_entropy")
    if direct:
        stable = _stable(direct, warmup_end)
        if not stable:
            return False, (
                f"No stable epochs for val_dir_entropy past warmup_end={warmup_end}"
            )
        m = mean(x["mean"] for x in stable)
        if m <= DIR_ENTROPY_FLOOR:
            return False, (
                f"val_dir_entropy mean = {m:.4f} <= {DIR_ENTROPY_FLOOR:.4f} "
                f"(0.8 × log 4)"
            )
        return True, (
            f"val_dir_entropy mean = {m:.4f} > {DIR_ENTROPY_FLOOR:.4f} "
            f"(0.8 × log 4)"
        )

    # Fallback: derive from per-direction distribution.
    dist_keys = [
        "val_dir_dist_short",
        "val_dir_dist_hold",
        "val_dir_dist_long",
        "val_dir_dist_flat",
    ]
    dist_series = [agg.get(k) for k in dist_keys]
    if all(s for s in dist_series):
        # Align by epoch on stable epochs.
        stables = [_stable(s, warmup_end) for s in dist_series]
        if any(not s for s in stables):
            return False, (
                f"No stable epochs for val_dir_dist_* past warmup_end={warmup_end}"
            )
        # Average per-direction mean across stable epochs, then compute entropy.
        means = [mean(x["mean"] for x in s) for s in stables]
        total = sum(means)
        if total <= 0:
            return False, "val_dir_dist_* sums to <= 0; cannot compute entropy"
        probs = [m / total for m in means]
        ent = _entropy_from_dist(probs)
        if ent <= DIR_ENTROPY_FLOOR:
            return False, (
                f"derived dir entropy = {ent:.4f} <= {DIR_ENTROPY_FLOOR:.4f}"
            )
        return True, (
            f"derived dir entropy = {ent:.4f} > {DIR_ENTROPY_FLOOR:.4f}"
        )

    return False, (
        "neither val_dir_entropy nor val_dir_dist_* present in aggregate — "
        "add a direction-distribution metric to HEALTH_DIAG to enable "
        "Tier 2 entropy check"
    )


def main():
    ap = argparse.ArgumentParser(
        description="Tier 2 behavioural exit check (Plan 5 Task 4).",
    )
    ap.add_argument(
        "metrics_json",
        help="Aggregated metrics JSON (output of aggregate-multi-seed-metrics.py).",
    )
    ap.add_argument(
        "--warmup-end",
        type=int,
        default=15,
        help="Epoch at which warmup ends (stable epochs are > this); default 15.",
    )
    args = ap.parse_args()

    try:
        with open(args.metrics_json) as fh:
            data = json.load(fh)
    except (OSError, json.JSONDecodeError) as exc:
        print(f"FAIL: cannot read {args.metrics_json}: {exc}")
        sys.exit(1)

    agg = data.get("aggregates")
    if not isinstance(agg, dict):
        print("FAIL: missing or malformed `aggregates` block in metrics JSON")
        sys.exit(1)

    checks = [
        check_trades_per_bar(agg, args.warmup_end),
        check_active_frac(agg, args.warmup_end),
        check_dir_entropy(agg, args.warmup_end),
    ]

    all_pass = all(ok for ok, _ in checks)
    for ok, msg in checks:
        print(("PASS" if ok else "FAIL") + ": " + msg)
    sys.exit(0 if all_pass else 1)


if __name__ == "__main__":
    main()