#!/usr/bin/env python3
"""Aggregate per-seed HEALTH_DIAG metrics into a single JSON report.

Plan 5 Task 1B.2 — consumes the concatenated logs from N*K (seed, fold)
training jobs (output of `argo logs --selector foxhunt-tag=<tag>`),
parses HEALTH_DIAG lines, groups by (seed, fold, epoch, metric_name), and
emits per-(epoch, metric) aggregate statistics across seeds.

HEALTH_DIAG line format (from training_loop.rs):

    HEALTH_DIAG[<epoch>]: <key1> [<sub1>=<val1> <sub2>=<val2> ...] \
                         <key2> [<sub3>=<val3> ...] ...

The line may be wrapped in tracing/JSON envelope:

    {"timestamp":..., "fields":{"message":"HEALTH_DIAG[N]: ..."}, ...}

Per-line followups (`reward_split`, `aux`) are also parsed.

Seed/fold attribution:

    The script reads `--multi-seed N --folds K` and assumes the input log
    contains exactly N*K independent training streams in order. Each stream
    is identified by the ordered (seed, fold) pair derived from its
    position in the log — Argo's `argo logs --no-color` interleaves pod
    output with a `pod_name` prefix that we use as the stream key. If
    pod-name attribution is unavailable, we fall back to detecting epoch
    resets (a HEALTH_DIAG[0] line marks a new stream).

Output JSON schema:

    {
      "tag": str,
      "multi_seed": int,
      "folds": int,
      "warmup_end_epoch": int | null,
      "streams_seen": int,
      "aggregates": {
        "<metric_name>": [
          {"epoch": 0, "mean": float, "std": float,
           "median": float, "min": float, "max": float,
           "n_samples": int},
          ...
        ],
        ...
      }
    }
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from collections import defaultdict
from pathlib import Path
from statistics import mean, median, pstdev
from typing import Iterable

# HEALTH_DIAG[<epoch>]: <body>
HEALTH_DIAG_RE = re.compile(r"HEALTH_DIAG\[(\d+)\]:\s*(.*)")

# JSON envelope: {"fields":{"message":"HEALTH_DIAG[..."}, ...}
JSON_MESSAGE_RE = re.compile(r'"message"\s*:\s*"([^"]*HEALTH_DIAG\[[^"]*)"')

# Top-level block: <name> [<key=val> <key=val> ...]
# Names are alphanumeric + underscore; the bracketed body contains key=val
# pairs separated by whitespace. We keep names simple (no nested brackets
# inside HEALTH_DIAG bodies — verified against current training_loop.rs).
BLOCK_RE = re.compile(r"(\w+)\s*\[([^\[\]]*)\]")

# key=val inside a block. Values may be:
#   - signed floats (with/without scientific notation): -3.911565e0
#   - signed ints: 5
#   - bool literals: true / false / on / off / ready / ...
#   - small enums (action_entropy=0.79, plasticity=ready)
# We capture as (key, value) and let downstream typing decide.
KV_RE = re.compile(r"(\w+)=([^\s\]]+)")

# Pod-name prefix from `argo logs` output (best-effort; not all log
# collectors emit this). Format: "<pod-name>: <log line>".
POD_PREFIX_RE = re.compile(r"^(?P<pod>[a-z0-9][a-z0-9-]*?):\s+(?P<rest>.*)$")


def parse_value(raw: str) -> float | None:
    """Parse a HEALTH_DIAG value. Returns None for non-numeric (boolean,
    enum) values — those are excluded from aggregation since mean/std are
    not meaningful for them."""
    # Strip trailing punctuation that sometimes appears at end of bracket.
    raw = raw.strip().rstrip(",;:")
    try:
        return float(raw)
    except ValueError:
        return None


def extract_health_diag(line: str) -> tuple[int, str] | None:
    """Pull (epoch, body) from a raw log line. Handles both bare
    HEALTH_DIAG and JSON-envelope-wrapped variants."""
    # JSON envelope first — common in production logs.
    m = JSON_MESSAGE_RE.search(line)
    if m:
        msg = m.group(1).replace("\\\"", '"')
        m2 = HEALTH_DIAG_RE.search(msg)
        if m2:
            return int(m2.group(1)), m2.group(2)
        return None
    # Bare HEALTH_DIAG.
    m = HEALTH_DIAG_RE.search(line)
    if m:
        return int(m.group(1)), m.group(2)
    return None


def parse_blocks(body: str) -> dict[str, float]:
    """Parse all `<block>[<kv> <kv> ...]` segments from a HEALTH_DIAG
    body. Returns flat dict with keys `<block>_<subkey>` (single underscore
    joiner — this is what the Plan 5 Task 4 tier check scripts expect, and
    what the synthetic test fixtures encode: e.g. `val [sharpe=…]` becomes
    aggregate key `val_sharpe`)."""
    out: dict[str, float] = {}
    for block_match in BLOCK_RE.finditer(body):
        block_name = block_match.group(1)
        block_body = block_match.group(2)
        for kv_match in KV_RE.finditer(block_body):
            key = kv_match.group(1)
            val = parse_value(kv_match.group(2))
            if val is not None:
                out[f"{block_name}_{key}"] = val
    return out


def stream_key(line: str, fallback: int) -> str:
    """Best-effort stream identifier for grouping HEALTH_DIAG lines back
    to a single training run. Uses pod-name prefix when present, else
    falls back to a synthetic counter that increments on epoch=0
    re-observations (see `attribute_streams`)."""
    m = POD_PREFIX_RE.match(line.strip())
    if m:
        return m.group("pod")
    return f"stream-{fallback}"


def attribute_streams(
    lines: Iterable[str],
    expected_streams: int,
) -> dict[str, list[tuple[int, dict[str, float]]]]:
    """Walk the log line-by-line and bucket HEALTH_DIAG observations by
    stream. Each stream's value is an ordered list of (epoch, metrics_dict)
    pairs.

    Stream attribution strategy:
      1. If pod-name prefix is present and stable, use it.
      2. Else, treat epoch=0 as a stream-boundary marker — each new
         HEALTH_DIAG[0] starts a new synthetic stream.
    """
    by_stream: dict[str, list[tuple[int, dict[str, float]]]] = defaultdict(list)
    synthetic_counter = 0
    current_synthetic = None
    last_epoch: int | None = None
    for raw in lines:
        diag = extract_health_diag(raw)
        if diag is None:
            continue
        epoch, body = diag
        metrics = parse_blocks(body)
        if not metrics:
            continue
        # Try pod-name attribution first.
        m = POD_PREFIX_RE.match(raw.strip())
        if m:
            key = m.group("pod")
        else:
            # Synthetic: a new stream starts when we either see the first
            # HEALTH_DIAG ever, or when the current epoch goes BACKWARDS
            # (epoch < last_epoch) — that signals a new training run.
            # Each HEALTH_DIAG epoch typically emits multiple sub-lines
            # (main / reward_split / aux); naively triggering on every
            # epoch=0 over-segments those sub-lines into separate streams.
            need_new_stream = current_synthetic is None or (
                last_epoch is not None and epoch < last_epoch
            )
            if need_new_stream:
                current_synthetic = f"stream-{synthetic_counter}"
                synthetic_counter += 1
            key = current_synthetic
        by_stream[key].append((epoch, metrics))
        last_epoch = epoch
    return dict(by_stream)


def aggregate_streams(
    by_stream: dict[str, list[tuple[int, dict[str, float]]]],
) -> dict[str, list[dict[str, float]]]:
    """Cross-stream aggregation. For each (epoch, metric_name) pair,
    compute mean / std / median / min / max across streams."""
    # (metric, epoch) → [values across streams]
    bucket: dict[tuple[str, int], list[float]] = defaultdict(list)
    for stream_metrics in by_stream.values():
        for epoch, metrics in stream_metrics:
            for metric_name, value in metrics.items():
                bucket[(metric_name, epoch)].append(value)

    # Group by metric name, sort by epoch.
    by_metric: dict[str, list[dict[str, float]]] = defaultdict(list)
    for (metric_name, epoch), values in bucket.items():
        n = len(values)
        if n == 0:
            continue
        by_metric[metric_name].append(
            {
                "epoch": epoch,
                "mean": mean(values),
                # pstdev (population std) so n=1 yields 0 instead of error;
                # documents that this is across-seed dispersion, not a
                # sample estimate.
                "std": pstdev(values) if n > 1 else 0.0,
                "median": median(values),
                "min": min(values),
                "max": max(values),
                "n_samples": n,
            }
        )
    for metric_name in by_metric:
        by_metric[metric_name].sort(key=lambda r: r["epoch"])
    return dict(by_metric)


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Aggregate HEALTH_DIAG metrics across multi-seed training runs.",
    )
    ap.add_argument(
        "--input",
        required=True,
        type=Path,
        help="Concatenated log file (output of argo logs --no-color)",
    )
    ap.add_argument(
        "--multi-seed",
        type=int,
        default=1,
        help="Number of seeds (informational; written to output JSON)",
    )
    ap.add_argument(
        "--folds",
        type=int,
        default=1,
        help="Number of folds (informational; written to output JSON)",
    )
    ap.add_argument(
        "--tag",
        type=str,
        default="",
        help="Workflow tag (informational; written to output JSON)",
    )
    ap.add_argument(
        "--warmup-end-epoch",
        type=int,
        default=None,
        help="Epoch at which warmup ends (informational; written to output JSON)",
    )
    ap.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Path for the aggregated JSON output",
    )
    args = ap.parse_args()

    if not args.input.is_file():
        print(f"ERROR: input file not found: {args.input}", file=sys.stderr)
        return 1

    expected_streams = args.multi_seed * args.folds
    with args.input.open("r", encoding="utf-8", errors="replace") as fh:
        by_stream = attribute_streams(fh, expected_streams)

    if not by_stream:
        print(
            f"ERROR: no HEALTH_DIAG lines found in {args.input}",
            file=sys.stderr,
        )
        return 1

    aggregates = aggregate_streams(by_stream)

    out = {
        "tag": args.tag,
        "multi_seed": args.multi_seed,
        "folds": args.folds,
        "warmup_end_epoch": args.warmup_end_epoch,
        "streams_seen": len(by_stream),
        "aggregates": aggregates,
    }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as fh:
        json.dump(out, fh, indent=2)

    print(
        f"Aggregated {len(by_stream)} streams "
        f"(expected {expected_streams}={args.multi_seed}x{args.folds}), "
        f"{len(aggregates)} unique metrics → {args.output}",
        file=sys.stderr,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())
