#!/usr/bin/env python3
"""Frozen Phase-5 bridge/control ledger analysis."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
from typing import Any

import numpy as np


SEED = 20260716
RESAMPLES = 100_000
RATE = 0.4725
ARMS = ("base", "A1", "A2", "A3", "A4")
MECHANISMS = ("A1", "A2", "A3", "A4")
CAPTURE_SIZES = {
    1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88,
    96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192,
    200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336,
    352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512,
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def numeric(values: list[float | int | None]) -> dict[str, Any]:
    finite = [float(value) for value in values if value is not None and math.isfinite(float(value))]
    return {
        "n": len(values),
        "finite_n": len(finite),
        "missing_n": len(values) - len(finite),
        "min": min(finite) if finite else None,
        "max": max(finite) if finite else None,
        "distinct_n": len(set(finite)),
    }


def ci(draws: np.ndarray) -> list[float]:
    low, high = np.quantile(draws, [0.025, 0.975])
    return [float(low), float(high)]


def cv(values: list[float]) -> float:
    array = np.asarray(values, dtype=np.float64)
    mean = float(array.mean())
    if mean == 0:
        return 0.0 if np.all(array == 0) else math.inf
    return float(array.std(ddof=0) / mean)


def parse_run(run_dir: Path) -> dict[str, Any]:
    result = json.loads((run_dir / "client/result.json").read_text())
    complete = json.loads((run_dir / "run-complete.json").read_text())
    t0 = int(result["t0_mono_ns"])
    clean_start = float(result["clean"]["start_s"])
    clean_end = float(result["clean"]["end_s"])
    stream = next((run_dir / "opprof").glob("*.jsonl"))
    records = []
    for line in stream.read_text().splitlines():
        item = json.loads(line)
        if "step_index" not in item:
            continue
        relative = (int(item["submit_mono_ns"]) - t0) / 1e9
        if clean_start <= relative < clean_end:
            item["relative_s"] = relative
            records.append(item)
    blocks = []
    decode_means = []
    waiting_means = []
    for block_index in range(48):
        start = clean_start + 5 * block_index
        selected = [item for item in records if start <= item["relative_s"] < start + 5]
        if not selected:
            # Fixed-time blocks are the resampling unit.  A rate-following trace
            # can legitimately have an idle block, which contributes no tokens
            # and no model-step time but must remain in the arrival-variance
            # distribution.
            blocks.append([0, 0.0])
            decode_means.append(0.0)
            waiting_means.append(0.0)
            continue
        tokens = sum(int(item["prefill_tokens"]) + int(item["decode_tokens"]) for item in selected)
        duration = sum(
            (int(item["complete_mono_ns"]) - int(item["submit_mono_ns"])) / 1e6
            for item in selected
        )
        blocks.append([tokens, duration])
        decode_means.append(float(np.mean([int(item["decode_batch_size"]) for item in selected])))
        waiting_means.append(float(np.mean([int(item["queues"]["waiting"]) for item in selected])))
    pure_decode = [
        item for item in records
        if int(item["prefill_tokens"]) == 0 and int(item["decode_batch_size"]) > 0
    ]
    pure_bucket = sum(int(item["cudagraph"]["bucket_tokens"]) for item in pure_decode)
    pure_padding = sum(int(item["cudagraph"]["padding_tokens"]) for item in pure_decode)
    support = sorted({int(item["decode_batch_size"]) for item in pure_decode})
    covered = sum(int(item["decode_batch_size"]) in CAPTURE_SIZES for item in pure_decode)
    prefix_queries = 0
    prefix_hits = 0
    prefix_present = 0
    for item in records:
        local = item.get("prefix", {}).get("local")
        if local is None:
            continue
        prefix_present += 1
        prefix_queries += int(local.get("queries", 0))
        prefix_hits += int(local.get("hits", 0))
    block_array = np.asarray(blocks, dtype=np.float64)
    return {
        "run_id": complete["run_id"],
        "run_dir": str(run_dir),
        "stream_sha256": sha256_file(stream),
        "blocks": block_array,
        "token_efficiency_per_ms": float(block_array[:, 0].sum() / block_array[:, 1].sum()),
        "clean_steps": len(records),
        "clean_duration_s": clean_end - clean_start,
        "clean_failed": int(result["clean"]["failed"]),
        "offered_rps": float(result["clean"]["offered_rps"]),
        "drain_seconds": float(result["drain_seconds"]),
        "warmup_completions": int(complete["client"]["warmup_completions"]),
        "warmup_gate_branch": complete["client"].get("warmup_gate_branch", "P3-pre-A-P5-1"),
        "warmup_stability": complete["client"].get("warmup_stability"),
        "cold_start_gate": complete["client"].get("cold_start_gate"),
        "decode_B_block_cv": cv(decode_means),
        "waiting_block_cv": cv(waiting_means),
        "pure_decode_steps": len(pure_decode),
        "pure_decode_support": support,
        "pure_decode_support_coverage": covered / len(pure_decode),
        "pure_decode_padding_tokens": pure_padding,
        "pure_decode_bucket_tokens": pure_bucket,
        "pure_decode_padding_fraction": pure_padding / pure_bucket,
        "prefix_present_steps": prefix_present,
        "prefix_queries": prefix_queries,
        "prefix_hits": prefix_hits,
        "prefix_hit_ratio": prefix_hits / prefix_queries if prefix_queries else 0.0,
        "layer1_invariants": complete["layer1"]["invariants"],
        "client_invariants": complete["client"]["invariants"],
        "server_invariants": complete["server_invariants"],
        "drain_quarantined": bool(complete["drain_quarantined"]),
    }


def hierarchical_draws(runs: list[dict[str, Any]], rng: np.random.Generator) -> np.ndarray:
    count = len(runs)
    output = np.empty(RESAMPLES, dtype=np.float64)
    for start in range(0, RESAMPLES, 5000):
        size = min(5000, RESAMPLES - start)
        token_sum = np.zeros(size, dtype=np.float64)
        duration_sum = np.zeros(size, dtype=np.float64)
        selected_runs = rng.integers(0, count, size=(size, count))
        for position in range(count):
            block_indices = rng.integers(0, 48, size=(size, 48))
            for run_index, run in enumerate(runs):
                mask = selected_runs[:, position] == run_index
                if not np.any(mask):
                    continue
                blocks = run["blocks"]
                choices = block_indices[mask]
                token_sum[mask] += blocks[choices, 0].sum(axis=1)
                duration_sum[mask] += blocks[choices, 1].sum(axis=1)
        output[start : start + size] = token_sum / duration_sum
    return output


def point_efficiency(runs: list[dict[str, Any]]) -> float:
    tokens = sum(float(run["blocks"][:, 0].sum()) for run in runs)
    duration = sum(float(run["blocks"][:, 1].sum()) for run in runs)
    return tokens / duration


def two_sided_p(draws: np.ndarray) -> float:
    return float(min(1.0, 2 * min(np.mean(draws <= 0), np.mean(draws >= 0))))


def holm(pvalues: dict[str, float]) -> dict[str, float]:
    ordered = sorted(pvalues, key=pvalues.get)
    adjusted: dict[str, float] = {}
    running = 0.0
    total = len(ordered)
    for rank, key in enumerate(ordered):
        running = max(running, (total - rank) * pvalues[key])
        adjusted[key] = min(1.0, running)
    return adjusted


def discover_primary(root: Path) -> dict[str, list[dict[str, Any]]]:
    result = {arm: [] for arm in ARMS}
    for marker in sorted((root / "primary").glob("*/moderate/run-complete.json")):
        if "background" in str(marker):
            continue
        parsed = parse_run(marker.parent)
        arm = parsed["run_id"].split("-r", 1)[0]
        if arm in result:
            result[arm].append(parsed)
    if any(len(result[arm]) != 3 for arm in ARMS):
        raise RuntimeError(f"primary replication mismatch: { {key:len(value) for key,value in result.items()} }")
    return result


def control_runs(root: Path, p3root: Path, pattern: str) -> tuple[list[dict[str, Any]], str]:
    fresh = sorted((root / "primary").glob(f"control-{pattern}-r*-C00/moderate/run-complete.json"))
    if fresh:
        if len(fresh) != 3:
            raise RuntimeError(f"partial fresh control set: {pattern}: {len(fresh)}")
        return [parse_run(path.parent) for path in fresh], "P5-rerun"
    return [parse_run(p3root / f"primary/{pattern}-C00/moderate")], "P3-reused"


def aggregate_aux(runs: list[dict[str, Any]], key: str) -> float:
    return float(np.mean([float(run[key]) for run in runs]))


def analyze(root: Path, p3root: Path, private: Path) -> dict[str, Any]:
    primary = discover_primary(root)
    p3_base = [parse_run(p3root / "primary/P10-C00/moderate")]
    controls = {}
    control_source = {}
    for pattern in ("P03", "P04"):
        controls[pattern], control_source[pattern] = control_runs(root, p3root, pattern)
    rng = np.random.default_rng(SEED)
    draws = {arm: hierarchical_draws(primary[arm], rng) for arm in ARMS}
    p3_base_draws = hierarchical_draws(p3_base, rng)
    control_draws = {pattern: hierarchical_draws(controls[pattern], rng) for pattern in controls}
    points = {arm: point_efficiency(primary[arm]) for arm in ARMS}
    p3_base_point = point_efficiency(p3_base)
    control_points = {pattern: point_efficiency(controls[pattern]) for pattern in controls}

    bridge_draws = draws["A3"] - p3_base_draws
    bridge_point = points["A3"] - p3_base_point
    bridge_ci = ci(bridge_draws)
    bridge = {
        "point_delta": bridge_point,
        "relative_abs_delta": abs(bridge_point) / p3_base_point,
        "ci95": bridge_ci,
        "within_3pct": abs(bridge_point) / p3_base_point <= 0.03,
        "ci_contains_zero": bridge_ci[0] <= 0 <= bridge_ci[1],
    }
    bridge["reuse_passed"] = bridge["within_3pct"] and bridge["ci_contains_zero"]

    deltas = {arm: draws[arm] - draws["base"] for arm in MECHANISMS}
    raw_p = {arm: two_sided_p(deltas[arm]) for arm in MECHANISMS}
    holm_p = holm(raw_p)
    manifests = {
        name: json.loads((private / f"P10-{name}.jsonl.summary.json").read_text())
        for name in ("base", "A1", "A3")
    }
    base_prefix = aggregate_aux(primary["base"], "prefix_hit_ratio")
    a1_prefix = aggregate_aux(primary["A1"], "prefix_hit_ratio")
    base_padding = aggregate_aux(primary["base"], "pure_decode_padding_fraction")
    a2_padding = aggregate_aux(primary["A2"], "pure_decode_padding_fraction")
    padding_reduction = (base_padding - a2_padding) / base_padding if base_padding > 0 else 0.0
    manipulations = {
        "A1": {
            "passed": (
                manifests["base"]["r16"] - manifests["A1"]["r16"] >= 0.15
                and (manifests["base"]["r16"] - manifests["A1"]["r16"]) / manifests["base"]["r16"] >= 0.20
                and manifests["A1"]["max_added_delay_seconds"] <= 64
                and abs(a1_prefix - base_prefix) <= 0.01
            ),
            "base_R16": manifests["base"]["r16"],
            "ablated_R16": manifests["A1"]["r16"],
            "prefix_hit_ratio_delta": a1_prefix - base_prefix,
            "max_added_delay_seconds": manifests["A1"]["max_added_delay_seconds"],
        },
        "A2": {
            "passed": aggregate_aux(primary["A2"], "pure_decode_support_coverage") >= 0.99 and padding_reduction >= 0.90,
            "support_coverage": aggregate_aux(primary["A2"], "pure_decode_support_coverage"),
            "base_padding_fraction": base_padding,
            "ablated_padding_fraction": a2_padding,
            "padding_reduction": padding_reduction,
            "observed_support": sorted({value for run in primary["A2"] for value in run["pure_decode_support"]}),
        },
        "A3": {
            "passed": (
                aggregate_aux(primary["A3"], "decode_B_block_cv") < aggregate_aux(primary["base"], "decode_B_block_cv")
                and aggregate_aux(primary["A3"], "waiting_block_cv") < aggregate_aux(primary["base"], "waiting_block_cv")
            ),
            "base_decode_B_cv": aggregate_aux(primary["base"], "decode_B_block_cv"),
            "ablated_decode_B_cv": aggregate_aux(primary["A3"], "decode_B_block_cv"),
            "base_waiting_cv": aggregate_aux(primary["base"], "waiting_block_cv"),
            "ablated_waiting_cv": aggregate_aux(primary["A3"], "waiting_block_cv"),
        },
        "A4": {
            "passed": sum(run["prefix_queries"] for run in primary["A4"]) == 0 and sum(run["prefix_hits"] for run in primary["A4"]) == 0,
            "prefix_queries": sum(run["prefix_queries"] for run in primary["A4"]),
            "prefix_hits": sum(run["prefix_hits"] for run in primary["A4"]),
        },
    }

    ledgers = {}
    for control in ("P03", "P04"):
        denominator = control_draws[control] - draws["base"]
        point_denominator = control_points[control] - points["base"]
        denominator_ci = ci(denominator)
        stable_denominator = bool(
            point_denominator > 0
            and denominator_ci[0] > 0
            and np.mean(denominator <= 0) <= 0.05
        )
        rows = {}
        share_draws = {}
        for arm in MECHANISMS:
            share_draws[arm] = deltas[arm] / denominator
            point = (points[arm] - points["base"]) / point_denominator
            interval = ci(share_draws[arm])
            reportable = stable_denominator and manipulations[arm]["passed"]
            rows[arm] = {
                "delta_E": points[arm] - points["base"],
                "delta_E_ci95": ci(deltas[arm]),
                "share": point if reportable else None,
                "share_ci95": interval if reportable else None,
                "diagnostic_share": point,
                "diagnostic_share_ci95": interval,
                "share_status": (
                    "REPORTABLE"
                    if reportable
                    else (
                        "N/A—unstable control denominator"
                        if not stable_denominator
                        else "N/A—manipulation failed"
                    )
                ),
                "p_two_sided": raw_p[arm],
                "p_holm": holm_p[arm],
                "manipulation_passed": manipulations[arm]["passed"],
            }
        residual_draws = 1.0 - sum(share_draws.values())
        residual_point = 1.0 - sum(row["diagnostic_share"] for row in rows.values())
        ledgers[control] = {
            "status": "EVALUABLE" if stable_denominator else "INCONCLUSIVE—unstable denominator",
            "control_source": control_source[control],
            "control_E": control_points[control],
            "base_E": points["base"],
            "gap_E": point_denominator,
            "gap_ci95": denominator_ci,
            "denominator_nonpositive_fraction": float(np.mean(denominator <= 0)),
            "stable_denominator": stable_denominator,
            "mechanisms": rows,
            "diagnostic_share_sum": sum(row["diagnostic_share"] for row in rows.values()),
            "residual_interaction": residual_point,
            "residual_interaction_ci95": ci(residual_draws),
            "residual_status": (
                "REPORTABLE"
                if stable_denominator and all(item["passed"] for item in manipulations.values())
                else "DIAGNOSTIC_ONLY—incomplete official share ledger"
            ),
        }

    dominance = {}
    for arm in MECHANISMS:
        per_control = {}
        for control in ("P03", "P04"):
            row = ledgers[control]["mechanisms"][arm]
            if not ledgers[control]["stable_denominator"] or not row["manipulation_passed"]:
                per_control[control] = "NOT_EVALUABLE"
                continue
            direction_ok = row["delta_E"] > 0 if arm != "A4" else True
            per_control[control] = (
                row["share"] >= 0.30
                and row["share_ci95"][0] > 0.15
                and row["p_holm"] < 0.05
                and direction_ok
                and row["manipulation_passed"]
            )
        dominance[arm] = {
            "per_control": per_control,
            "verdict": (
                "NOT EVALUABLE"
                if any(value == "NOT_EVALUABLE" for value in per_control.values())
                else (
                    "DOMINANT"
                    if all(per_control.values())
                    else ("CONTROL-SENSITIVE" if any(per_control.values()) else "NOT DOMINANT")
                )
            ),
        }

    all_runs = [run for arm in ARMS for run in primary[arm]]
    share_widths = [
        ledgers[control]["mechanisms"][arm]["diagnostic_share_ci95"][1]
        - ledgers[control]["mechanisms"][arm]["diagnostic_share_ci95"][0]
        for control in ("P03", "P04") for arm in MECHANISMS
    ]
    state = json.loads((root / "controller-state.json").read_text())
    amendment_evidence_path = root / "a-p5-1-retained-audit.jsonl"
    amendment_evidence = []
    if amendment_evidence_path.exists():
        amendment_evidence = [
            json.loads(line) for line in amendment_evidence_path.read_text().splitlines()
        ]
    invariants = {
        "primary_runs_15": len(all_runs) == 15,
        "three_replicates_per_arm": all(len(primary[arm]) == 3 for arm in ARMS),
        "clean_duration_240": all(math.isclose(run["clean_duration_s"], 240.0) for run in all_runs),
        "clean_failures_zero": all(run["clean_failed"] == 0 for run in all_runs),
        "offered_rate_within_5pct": all(abs(run["offered_rps"] / RATE - 1) <= 0.05 for run in all_runs),
        "layer1_accounting": all(all(run["layer1_invariants"].values()) for run in all_runs),
        "client_invariants": all(all(value for key, value in run["client_invariants"].items() if key != "drain_re_adjudicated") for run in all_runs),
        "server_invariants": all(all(run["server_invariants"].values()) for run in all_runs),
        "a_p5_1_cold_start_gates": all(
            run["cold_start_gate"] is not None
            and run["cold_start_gate"]["passed"]
            for run in all_runs
        ),
        "drain_quarantine_under_20pct": sum(run["drain_quarantined"] for run in all_runs) / 15 <= 0.20,
        "gpu_budget_below_6": float(state["gpu_hours_total"]) < 6.0,
        "manifests_same_ids": len({manifests[name]["request_id_set_sha256"] for name in manifests}) == 1,
        "manifests_same_token_sums": len({manifests[name]["input_tokens"]["sum"] for name in manifests}) == 1 and len({manifests[name]["output_tokens"]["sum"] for name in manifests}) == 1,
        "control_denominators_stable": all(ledgers[control]["stable_denominator"] for control in ledgers),
        "bridge_decision_resolved": bridge["reuse_passed"] or all(source == "P5-rerun" for source in control_source.values()),
        "ratios_finite": all(math.isfinite(ledgers[c]["mechanisms"][a]["diagnostic_share"]) for c in ledgers for a in MECHANISMS),
        "per_arm_results_not_all_identical": len({round(points[arm], 12) for arm in ARMS}) > 1,
    }
    red_flags = [key for key, value in invariants.items() if not value]
    publishable = (
        not red_flags
        and all(item["passed"] for item in manipulations.values())
        and sum(width > 0.50 for width in share_widths) < 2
    )
    return {
        "schema": 1,
        "status": "PUBLISHABLE" if publishable else "INCONCLUSIVE_OR_PARTIAL",
        "limitation": "Recorded-arrival P5 bridge ledger anchored to P3 controls; not a literal decomposition of P3's already-uniform P10 gap.",
        "analysis_seed": SEED,
        "bootstrap_resamples": RESAMPLES,
        "efficiency": {
            arm: {
                "point": points[arm],
                "ci95": ci(draws[arm]),
                "runs": [
                    {key: value for key, value in run.items() if key not in {"blocks", "warmup_stability"}}
                    for run in primary[arm]
                ],
            }
            for arm in ARMS
        },
        "p3_base_E": p3_base_point,
        "bridge": bridge,
        "control_sources": control_source,
        "manipulations": manipulations,
        "holm": {"family": list(MECHANISMS), "raw_p": raw_p, "adjusted_p": holm_p},
        "ledgers": ledgers,
        "dominance": dominance,
        "config_tier_A2": {
            "delta_E": points["A2"] - points["base"],
            "relative_E_delta": points["A2"] / points["base"] - 1,
            "delta_E_ci95": ci(deltas["A2"]),
            "base_padding_fraction": base_padding,
            "A2_padding_fraction": a2_padding,
            "padding_reduction": padding_reduction,
        },
        "amendment_A_P5_1": {
            "reason": "Rate-following throughput drift tracks arrival shape and is not a cold-start stationarity test.",
            "retained_failed_run_evidence": amendment_evidence,
            "recorded_drift_range": (
                [
                    min(
                        item["superseded_drift_evidence"]["normalized_drift"]
                        for item in amendment_evidence
                        if item["run"] != "A3-r1-C00"
                    ),
                    max(
                        item["superseded_drift_evidence"]["normalized_drift"]
                        for item in amendment_evidence
                        if item["run"] != "A3-r1-C00"
                    ),
                ]
                if amendment_evidence else None
            ),
            "uniform_A3_drift": (
                next(
                    item["superseded_drift_evidence"]["normalized_drift"]
                    for item in amendment_evidence
                    if item["run"] == "A3-r1-C00"
                )
                if amendment_evidence else None
            ),
        },
        "gpu": {
            "new_h20_hours": float(state["gpu_hours_total"]),
            "hard_cap": 6.0,
            "controller_status": state["status"],
            "completed_measured_runs_including_background": state["completed_measured_runs"],
            "completed_burnins": state["completed_burnins"],
        },
        "sanity": {
            "red_flags": red_flags,
            "invariants": invariants,
            "numeric": {
                "primary_E": numeric([run["token_efficiency_per_ms"] for run in all_runs]),
                "clean_steps": numeric([run["clean_steps"] for run in all_runs]),
                "offered_rps": numeric([run["offered_rps"] for run in all_runs]),
                "drain_seconds": numeric([run["drain_seconds"] for run in all_runs]),
                "diagnostic_share": numeric([ledgers[c]["mechanisms"][a]["diagnostic_share"] for c in ledgers for a in MECHANISMS]),
                "residual_interaction": numeric([ledgers[c]["residual_interaction"] for c in ledgers]),
                "share_ci_width": numeric(share_widths),
            },
            "declared": {
                "manipulation_failures": [arm for arm, item in manipulations.items() if not item["passed"]],
                "control_sources": control_source,
                "bridge_reuse_passed": bridge["reuse_passed"],
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--p3-root", type=Path, required=True)
    parser.add_argument("--private", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    result = analyze(args.root, args.p3_root, args.private)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, sort_keys=True, indent=2) + "\n")
    print(json.dumps({
        "status": result["status"],
        "bridge": result["bridge"],
        "red_flags": result["sanity"]["red_flags"],
        "gpu": result["gpu"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()
