#!/usr/bin/env python3
"""Plan 5 Task 2 — populate `config/metric-bands.toml` from N known-good aggregate JSONs.

Consumes the per-stream aggregate JSONs produced by
`scripts/aggregate-multi-seed-metrics.py` (Plan 5 Task 1B.2) and emits TOML
band entries at `mean ± 5σ` (warn) / `mean ± 10σ` (error). The N>=2 case is the
production path; for N=1 we fall back to safe symmetric multiplicative
defaults documented inline.

Usage:

    populate-metric-bands-from-runs.py \\
        --aggregate run-a.json --aggregate run-b.json --aggregate run-c.json \\
        --warmup-end-epoch 10 \\
        --include-metrics avg_q_value train_loss val_sharpe ... \\
        --output config/metric-bands.toml

Each `--aggregate` JSON is expected to contain the schema written by
`aggregate-multi-seed-metrics.py`:

    { "aggregates": { "<metric>": [ {"epoch": int, "mean": float, ...}, ... ], ... } }

The script computes per-metric POST-WARMUP statistics (epoch >= warmup-end)
across the union of all samples from all aggregate JSONs:

    band_warn  = [grand_mean - 5σ,  grand_mean + 5σ]
    band_error = [grand_mean - 10σ, grand_mean + 10σ]

Where σ is the population standard deviation of the per-(epoch, run) sample
values (NOT the std of per-run means — we want to capture WITHIN-RUN epoch
drift in the band, not just CROSS-RUN dispersion).

For metrics where the variance is genuinely zero post-warmup (e.g.
`tau_eff = 0.005` always), we widen the band slightly so a single out-of-
band epoch doesn't trip on numerical noise:

    σ_eff = max(σ, |grand_mean| * 1e-4 + 1e-9)

For N=1 sample we fall back to multiplicative defaults — same shape as the
plan's `[settings]` block but multiplicative since we have no sigma:

    band_warn  = [|x| * 0.5,  |x| * 2.0]   (signed: [min(x*0.5,x*2), max(x*0.5,x*2)])
    band_error = [|x| * 0.1,  |x| * 10.0]

Existing bands in the output toml are PRESERVED unless --force is passed.
The `[settings]` block is preserved verbatim.
"""
from __future__ import annotations

import argparse
import json
import math
import sys
from collections import defaultdict
from pathlib import Path
from statistics import pstdev


def load_aggregate_json(path: Path) -> dict:
    """Load an aggregate JSON written by aggregate-multi-seed-metrics.py."""
    with path.open("r", encoding="utf-8") as fh:
        return json.load(fh)


def collect_post_warmup_samples(
    aggregates: list[dict],
    warmup_end_epoch: int,
) -> dict[str, list[float]]:
    """For each metric, collect per-epoch samples from epoch >= warmup_end across all runs.

    Each aggregate's `aggregates.<metric>[i].mean` already represents the
    cross-seed mean at epoch i. We keep that as a single sample (one per
    run × epoch). Cross-run dispersion is captured by including all runs'
    means; within-run epoch drift is captured by including all post-warmup
    epochs.
    """
    by_metric: dict[str, list[float]] = defaultdict(list)
    for agg in aggregates:
        for metric_name, rows in agg.get("aggregates", {}).items():
            for row in rows:
                if row.get("epoch", -1) >= warmup_end_epoch:
                    val = row.get("mean")
                    if val is not None and math.isfinite(val):
                        by_metric[metric_name].append(float(val))
    return dict(by_metric)


def compute_bands(
    samples: list[float],
    warn_sigma: float = 5.0,
    error_sigma: float = 10.0,
) -> dict[str, float] | None:
    """Compute warn / error bands from N samples. Returns None if N==0."""
    n = len(samples)
    if n == 0:
        return None
    mean_val = sum(samples) / n
    if n == 1:
        # Multiplicative fallback — captures sign while widening generously.
        x = samples[0]
        # Symmetric multiplicative shape; collapse to additive for x near 0.
        if abs(x) < 1e-9:
            return {
                "warn_low": -1.0, "warn_high": 1.0,
                "error_low": -10.0, "error_high": 10.0,
            }
        warn_lo = min(x * 0.5, x * 2.0)
        warn_hi = max(x * 0.5, x * 2.0)
        err_lo = min(x * 0.1, x * 10.0)
        err_hi = max(x * 0.1, x * 10.0)
        # Make sure error envelope strictly contains warn envelope.
        err_lo = min(err_lo, warn_lo - abs(warn_hi - warn_lo))
        err_hi = max(err_hi, warn_hi + abs(warn_hi - warn_lo))
        return {
            "warn_low": warn_lo, "warn_high": warn_hi,
            "error_low": err_lo, "error_high": err_hi,
        }
    sigma = pstdev(samples)
    # Effective sigma — never zero.
    sigma_eff = max(sigma, abs(mean_val) * 1e-4 + 1e-9)
    return {
        "warn_low": mean_val - warn_sigma * sigma_eff,
        "warn_high": mean_val + warn_sigma * sigma_eff,
        "error_low": mean_val - error_sigma * sigma_eff,
        "error_high": mean_val + error_sigma * sigma_eff,
    }


def format_toml(
    settings_block: str,
    bands: dict[str, dict[str, float]],
) -> str:
    """Render a metric-bands.toml document with the given settings + bands.

    Settings block is taken verbatim from the existing file (or a sensible
    default). Bands are written under a flat `[bands]` table with one
    sub-table per metric:

        [bands.<metric>]
        warn_low = ...
        warn_high = ...
        error_low = ...
        error_high = ...
    """
    lines: list[str] = []
    lines.append("# Regression-detection bands for DQN training metrics.")
    lines.append("# Invariant 7. Per A.4, bands populated from last 3 known-good runs.")
    lines.append("# Band = [mean - 5σ, mean + 5σ] per metric (warn);")
    lines.append("#        [mean - 10σ, mean + 10σ] (error). Recomputed after baseline promotion.")
    lines.append("# N consecutive out-of-band epochs → WARN; 2N consecutive → ERROR + clean termination.")
    lines.append("")
    lines.append(settings_block.rstrip())
    lines.append("")
    lines.append("[bands]")
    lines.append("")
    for metric in sorted(bands.keys()):
        b = bands[metric]
        lines.append(f"[bands.{metric}]")
        lines.append(f"warn_low  = {_fmt_float(b['warn_low'])}")
        lines.append(f"warn_high = {_fmt_float(b['warn_high'])}")
        lines.append(f"error_low  = {_fmt_float(b['error_low'])}")
        lines.append(f"error_high = {_fmt_float(b['error_high'])}")
        lines.append("")
    return "\n".join(lines).rstrip() + "\n"


def _fmt_float(x: float) -> str:
    """TOML-friendly float rendering with reasonable precision."""
    if not math.isfinite(x):
        # toml has no NaN/Inf; use a sentinel that downstream load-time
        # validation will reject loudly.
        return "0.0  # WARNING: non-finite computed band — please check inputs"
    if x == 0.0:
        return "0.0"
    if abs(x) >= 1e6 or abs(x) < 1e-3:
        return f"{x:.6e}"
    return f"{x:.6f}"


def parse_existing_settings(path: Path) -> str:
    """Extract the `[settings]` block from the existing toml (verbatim)."""
    if not path.is_file():
        return (
            "[settings]\n"
            "consecutive_epochs_for_warn = 3\n"
            "consecutive_epochs_for_error = 6\n"
        )
    text = path.read_text(encoding="utf-8")
    out_lines: list[str] = []
    in_block = False
    for line in text.splitlines():
        stripped = line.strip()
        if stripped == "[settings]":
            in_block = True
            out_lines.append(line)
            continue
        if in_block:
            if stripped.startswith("[") and stripped.endswith("]"):
                break
            out_lines.append(line)
    if not out_lines:
        return (
            "[settings]\n"
            "consecutive_epochs_for_warn = 3\n"
            "consecutive_epochs_for_error = 6\n"
        )
    return "\n".join(out_lines) + "\n"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument(
        "--aggregate",
        action="append",
        type=Path,
        default=[],
        help="Path to an aggregate JSON (output of aggregate-multi-seed-metrics.py). "
             "Pass once per known-good run. N>=3 is the plan-prescribed minimum; N=1 falls back to safe defaults.",
    )
    ap.add_argument(
        "--warmup-end-epoch",
        type=int,
        default=10,
        help="Epoch at which warmup ends — only epoch >= this contributes to band statistics. Default 10.",
    )
    ap.add_argument(
        "--include-metrics",
        nargs="+",
        default=[],
        help="If non-empty, restrict band emission to these metric names only. Otherwise emit all metrics found.",
    )
    ap.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output TOML path (typically config/metric-bands.toml).",
    )
    ap.add_argument(
        "--force",
        action="store_true",
        help="Overwrite existing band entries instead of preserving them. Settings block is always preserved.",
    )
    args = ap.parse_args()

    if not args.aggregate:
        print("ERROR: at least one --aggregate JSON is required", file=sys.stderr)
        return 1

    aggregates = []
    for p in args.aggregate:
        if not p.is_file():
            print(f"ERROR: aggregate JSON not found: {p}", file=sys.stderr)
            return 1
        aggregates.append(load_aggregate_json(p))

    by_metric = collect_post_warmup_samples(aggregates, args.warmup_end_epoch)

    if args.include_metrics:
        keep = set(args.include_metrics)
        by_metric = {k: v for k, v in by_metric.items() if k in keep}

    bands: dict[str, dict[str, float]] = {}
    for metric_name, samples in by_metric.items():
        b = compute_bands(samples)
        if b is not None:
            bands[metric_name] = b

    if not bands:
        print(
            "ERROR: no bands computable — check warmup-end-epoch and metric names "
            "(post-warmup samples may be empty)",
            file=sys.stderr,
        )
        return 1

    settings_block = parse_existing_settings(args.output)

    # Preserve existing bands unless --force.
    if not args.force and args.output.is_file():
        existing_text = args.output.read_text(encoding="utf-8")
        # Naive parse — we look for `[bands.<name>]` headers and keep the
        # block intact if we're not generating a new one for that metric.
        # Full parse is overkill since we control the writer too.
        existing_blocks = _parse_band_blocks(existing_text)
        for metric, body in existing_blocks.items():
            if metric not in bands:
                bands[metric] = body

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(format_toml(settings_block, bands), encoding="utf-8")

    print(
        f"Wrote {len(bands)} band entries to {args.output} "
        f"(from {len(args.aggregate)} aggregate JSON(s), "
        f"warmup-end={args.warmup_end_epoch})",
        file=sys.stderr,
    )
    return 0


def _parse_band_blocks(text: str) -> dict[str, dict[str, float]]:
    """Naive parser for `[bands.<metric>]` sub-tables. Returns dict by metric name."""
    out: dict[str, dict[str, float]] = {}
    current_metric: str | None = None
    current_body: dict[str, float] = {}
    for line in text.splitlines():
        stripped = line.strip()
        if stripped.startswith("[bands.") and stripped.endswith("]"):
            if current_metric is not None and current_body:
                out[current_metric] = current_body
            current_metric = stripped[len("[bands."):-1]
            current_body = {}
            continue
        if current_metric is not None and "=" in stripped and not stripped.startswith("#"):
            key, _, value = stripped.partition("=")
            key = key.strip()
            value = value.split("#", 1)[0].strip()
            try:
                current_body[key] = float(value)
            except ValueError:
                pass
    if current_metric is not None and current_body:
        out[current_metric] = current_body
    return out


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