#!/usr/bin/env python3
"""Frozen Phase-6 paired-anchor, frontier, rank, and Layer-1 analysis."""

from __future__ import annotations

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

import numpy as np


OLD_BEST = "tp2_mns32"
TOP20 = {"tp2_mns32", "tp2_mns64", "tp4_mns16", "tp4_mns32", "tp4_mns64"}


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


def numeric(values: list[float | int | None]) -> dict[str, Any]:
    finite = [float(x) for x in values if x is not None and math.isfinite(float(x))]
    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 cv(values: list[float]) -> float:
    if not values:
        return 0.0
    array = np.asarray(values, dtype=np.float64)
    mean = float(array.mean())
    return float(array.std(ddof=0) / mean) if mean else 0.0


def floor_buckets(scores: dict[str, float]) -> tuple[float, dict[str, int]]:
    tol = max(1e-9, 1e-6 * max(abs(x) for x in scores.values()))
    return tol, {key: math.floor(value / tol) for key, value in scores.items()}


def kendall_tau_b(x: dict[str, int], y: dict[str, int]) -> dict[str, Any]:
    keys = sorted(x)
    c = d = tx = ty = joint = 0
    for i, left in enumerate(keys):
        for right in keys[i + 1:]:
            dx = (x[left] > x[right]) - (x[left] < x[right])
            dy = (y[left] > y[right]) - (y[left] < y[right])
            if dx == 0 and dy == 0:
                joint += 1
            elif dx == 0:
                tx += 1
            elif dy == 0:
                ty += 1
            elif dx == dy:
                c += 1
            else:
                d += 1
    denom = math.sqrt((c + d + tx) * (c + d + ty))
    return {
        "tau_b": (c - d) / denom if denom else None, "concordant": c,
        "discordant": d, "v20_only_ties": tx, "v24_only_ties": ty,
        "joint_ties": joint, "pairs": len(keys) * (len(keys) - 1) // 2,
    }


def layer_metrics(records: list[dict[str, Any]]) -> dict[str, Any]:
    model = [x for x in records if x.get("model_executed")]
    mode = Counter(str(x["cudagraph"]["runtime_mode"]) for x in model)
    total_bucket = sum(int(x["cudagraph"]["bucket_tokens"]) for x in model)
    total_padding = sum(int(x["cudagraph"]["padding_tokens"]) for x in model)
    kv = [float(x["kv"]["usage"]) for x in model]
    waiting = [float(x["queues"]["waiting"]) for x in model]
    decode_b = [float(x["decode_batch_size"]) for x in model]
    return {
        "steps": len(records), "model_steps": len(model),
        "prefill_only_steps": sum(int(x["prefill_tokens"]) > 0 and int(x["decode_tokens"]) == 0 for x in model),
        "decode_only_steps": sum(int(x["prefill_tokens"]) == 0 and int(x["decode_tokens"]) > 0 for x in model),
        "mixed_steps": sum(int(x["prefill_tokens"]) > 0 and int(x["decode_tokens"]) > 0 for x in model),
        "prefill_tokens": sum(int(x["prefill_tokens"]) for x in model),
        "decode_tokens": sum(int(x["decode_tokens"]) for x in model),
        "preemptions": sum(int(x["preemptions"]) for x in model),
        "kv_usage_mean": float(np.mean(kv)) if kv else None,
        "kv_usage_max": max(kv) if kv else None,
        "waiting_mean": float(np.mean(waiting)) if waiting else None,
        "waiting_max": max(waiting) if waiting else None,
        "waiting_cv": cv(waiting), "decode_B_mean": float(np.mean(decode_b)) if decode_b else None,
        "decode_B_cv": cv(decode_b), "graph_mode_counts": dict(sorted(mode.items())),
        "graph_mode_shares": {key: value / len(model) for key, value in sorted(mode.items())} if model else {},
        "padding_fraction": total_padding / total_bucket if total_bucket else 0.0,
    }


def load_stream(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text().splitlines() if "step_index" in json.loads(line)]


def p3_reference(p3_root: Path) -> dict[str, Any]:
    run = p3_root / "primary/P10-C00/moderate"
    result = json.loads((run / "client/result.json").read_text())
    t0 = int(result["t0_mono_ns"])
    lo = t0 + int(float(result["clean"]["start_s"]) * 1e9)
    hi = t0 + int(float(result["clean"]["end_s"]) * 1e9)
    stream = next((run / "opprof").glob("*.jsonl"))
    records = [x for x in load_stream(stream) if lo <= int(x["submit_mono_ns"]) < hi]
    return {
        "limitation": "P3 P10 reference is TP1/MNS-default, steady arrival, sampling_u<=0.125, output<=256; it is contextual rather than a matched v0.20 composition baseline.",
        "stream_sha256": sha256_file(stream), "metrics": layer_metrics(records),
    }


def accepted_anchor(anchor_dir: Path) -> dict[str, Any]:
    primary = json.loads((anchor_dir / "result.json").read_text())
    cell_dir = anchor_dir.parent
    anchor = float(primary["anchor"])
    confirms = []
    for path in sorted(cell_dir.glob("confirm-*-anchor-*/result.json")):
        item = json.loads(path.read_text())
        if math.isclose(float(item["anchor"]), anchor, rel_tol=0, abs_tol=1e-15):
            confirms.append(item)
    trials = [primary, *confirms]
    votes = [bool(x["feasible"]) for x in trials]
    if len(votes) == 1 or len(set(votes)) == 1:
        verdict = votes[0]
        resolved = True
    elif len(votes) >= 3:
        verdict = sum(votes) >= 2
        resolved = True
    else:
        verdict = None
        resolved = False
    return {
        "anchor": anchor, "primary": primary, "confirmations": confirms,
        "trial_count": len(trials), "pass_rates": [x["pass_rate"] for x in trials],
        "feasible_votes": votes, "accepted_feasible": verdict, "resolved": resolved,
        "accepted_pass_rate": float(np.median([x["pass_rate"] for x in trials])),
    }


def matching_anchor_dir(cell_dir: Path, anchor: float) -> Path | None:
    for path in cell_dir.glob("anchor-*/result.json"):
        item = json.loads(path.read_text())
        if math.isclose(float(item["anchor"]), anchor, rel_tol=0, abs_tol=1e-15):
            return path.parent
    return None


def analyze(root: Path, ground_path: Path, p3_root: Path) -> dict[str, Any]:
    ground = json.loads(ground_path.read_text())
    old_cells = {x["cell_id"]: x for x in ground["cells"]}
    old_probe = {
        (x["cell_id"], float(p["sampling_u"])): p
        for x in ground["cells"] for p in x["probe_history"]
    }
    cells: dict[str, Any] = {}
    all_primary = []
    all_client_invariants = []
    selection_matches = []
    for cell, old in old_cells.items():
        cell_dir = root / "solo-authoritative/cells" / cell
        colocated_dir = root / "cells" / cell
        f20 = float(old["measured_objective"]["slo_feasible_req_s_per_gpu"])
        if not (cell_dir / "cell-valid.json").exists():
            cells[cell] = {
                "tp": old["tensor_parallel_size"], "mns": old["max_num_seqs"],
                "measurement_status": "UNMEASURED_SOLO",
                "f20": f20, "f24": None, "drift": None,
                "observed_max_feasible_rate": None,
                "material_frontier_moved": None, "boundary": "UNMEASURED",
                "bounded": False, "censor": "UNMEASURED_SOLO",
                "anchors": [], "cell_valid": None,
            }
            continue
        valid = json.loads((cell_dir / "cell-valid.json").read_text())
        stream = next((cell_dir / "opprof").glob("*.jsonl"))
        stream_records = load_stream(stream)
        anchors = []
        for path in sorted(cell_dir.glob("anchor-*/result.json")):
            accepted = accepted_anchor(path.parent)
            primary = accepted["primary"]
            key = (cell, float(primary["anchor"]))
            historical = old_probe[key]
            lo = int(primary["interval"]["start_mono_ns"])
            hi = int(primary["interval"]["end_mono_ns"])
            records = [x for x in stream_records if lo <= int(x["submit_mono_ns"]) <= hi]
            accepted["layer1"] = layer_metrics(records)
            for confirmation in accepted["confirmations"]:
                confirm_lo = int(confirmation["interval"]["start_mono_ns"])
                confirm_hi = int(confirmation["interval"]["end_mono_ns"])
                confirmation["layer1"] = layer_metrics([
                    x for x in stream_records
                    if confirm_lo <= int(x["submit_mono_ns"]) <= confirm_hi
                ])
            accepted["v20"] = {
                "pass_rate": historical["pass_rate"], "feasible": historical["feasible"],
                "request_count": historical["request_count"],
                "rate_per_gpu": historical["request_rate_per_gpu_req_s_gpu"],
            }
            accepted["v24"] = {
                "pass_rate": accepted["accepted_pass_rate"],
                "feasible": accepted["accepted_feasible"],
                "request_count": primary["selection"]["count"],
                "rate_per_gpu": primary["selection"]["offered_req_s_per_gpu"],
            }
            colocated_anchor_dir = matching_anchor_dir(colocated_dir, float(primary["anchor"]))
            if colocated_anchor_dir is not None:
                colocated = accepted_anchor(colocated_anchor_dir)
                accepted["colocated"] = {
                    "primary_pass_rate": colocated["primary"]["pass_rate"],
                    "primary_feasible": colocated["primary"]["feasible"],
                    "trial_pass_rates": colocated["pass_rates"],
                    "accepted_feasible": colocated["accepted_feasible"],
                    "solo_minus_colocated_primary_pass_rate": (
                        accepted["accepted_pass_rate"] - colocated["primary"]["pass_rate"]
                    ),
                }
            else:
                accepted["colocated"] = None
            accepted["feasibility_flip"] = (
                accepted["accepted_feasible"] is not None
                and accepted["accepted_feasible"] != historical["feasible"]
            )
            anchors.append(accepted)
            all_primary.append(primary)
            all_client_invariants.extend(primary["invariants"].values())
            selection_matches.append(primary["selection"]["count"] == historical["request_count"])
        peak_u = float(old["measured_objective"]["best_sampling_u"])
        peak = next(x for x in anchors if math.isclose(x["anchor"], peak_u, abs_tol=1e-15))
        feasible = [x for x in anchors if x["accepted_feasible"] is True]
        infeasible = [x for x in anchors if x["accepted_feasible"] is False]
        unresolved = [x for x in anchors if x["accepted_feasible"] is None]
        observed_frontier = max((x["v24"]["rate_per_gpu"] for x in feasible), default=None)
        lower = [x for x in anchors if x["anchor"] < peak_u]
        upper = [x for x in anchors if x["anchor"] > peak_u]
        if unresolved:
            bounded, censor = False, "UNRESOLVED_SOLO_ANCHOR"
        elif feasible and infeasible:
            monotonic = max(x["anchor"] for x in feasible) < min(x["anchor"] for x in infeasible)
            bounded = monotonic
            censor = None if bounded else "NONMONOTONIC_SOLO_ANCHORS"
        elif feasible:
            bounded, censor = False, "RIGHT_CENSORED_HISTORY_EDGE"
        elif infeasible:
            bounded, censor = False, "LEFT_CENSORED_HISTORY_EDGE"
        else:
            bounded, censor = False, "NO_RESOLVED_SOLO_ANCHOR"
        frontier = (
            None if censor in {
                "UNRESOLVED_SOLO_ANCHOR", "NONMONOTONIC_SOLO_ANCHORS",
                "NO_RESOLVED_SOLO_ANCHOR",
            } else observed_frontier
        )
        drift = (frontier / f20 - 1) if frontier is not None else None
        if censor == "NONMONOTONIC_SOLO_ANCHORS":
            boundary = "NONMONOTONIC"
        elif peak["accepted_feasible"] is False:
            boundary = "DOWN"
        elif peak["accepted_feasible"] is None:
            boundary = "UNRESOLVED"
        elif any(x["accepted_feasible"] is True for x in upper):
            boundary = "UP"
        else:
            boundary = "STABLE"
        cells[cell] = {
            "tp": old["tensor_parallel_size"], "mns": old["max_num_seqs"],
            "measurement_status": "SOLO_AUTHORITATIVE",
            "f20": f20, "f24": frontier, "drift": drift,
            "observed_max_feasible_rate": observed_frontier,
            "material_frontier_moved": bounded and drift is not None and abs(drift) > .05,
            "boundary": boundary, "bounded": bounded, "censor": censor,
            "anchors": anchors, "cell_valid": valid,
        }

    scores20 = {key: value["f20"] for key, value in cells.items()}
    scores24 = {key: value["f24"] for key, value in cells.items() if value["f24"] is not None}
    tol20, buckets20 = floor_buckets(scores20)
    tol24, buckets24 = floor_buckets(scores24)
    full_bounded = len(scores24) == 12 and all(x["bounded"] for x in cells.values())
    max24 = max(buckets24.values())
    if not full_bounded:
        argmax = "INCONCLUSIVE"
    else:
        argmax = "SURVIVED" if buckets24[OLD_BEST] == max24 else "MOVED"
    tau = kendall_tau_b(buckets20, buckets24) if full_bounded else None
    reversals = []
    for left in sorted(TOP20):
        for right in sorted(TOP20):
            if left >= right:
                continue
            if not cells[left]["bounded"] or not cells[right]["bounded"]:
                continue
            if left not in scores24 or right not in scores24:
                continue
            old_delta = scores20[left] - scores20[right]
            if abs(old_delta) / max(scores20[left], scores20[right]) <= .05:
                continue
            new_delta = scores24[left] - scores24[right]
            if old_delta * new_delta < 0:
                reversals.append([left, right])
    if not full_bounded or argmax == "INCONCLUSIVE" or tau is None:
        ranking = "INCONCLUSIVE"
    elif argmax == "MOVED" or tau["tau_b"] < .8 or reversals:
        ranking = "MOVED"
    else:
        ranking = "SURVIVED"
    trap_inputs = ["tp4_mns8", "tp4_mns16", "tp4_mns32"]
    if not all(cells[x]["bounded"] for x in trap_inputs):
        trap = "INCONCLUSIVE"
    elif buckets24["tp4_mns16"] == max24:
        trap = "CEASES_TO_BE_A_TRAP"
    elif buckets24["tp4_mns16"] >= max(buckets24["tp4_mns8"], buckets24["tp4_mns32"]):
        trap = "PERSISTS"
    else:
        trap = "ESCAPES"

    colocated_state = json.loads((root / "controller-state.json").read_text())
    solo_state_path = root / "solo-authoritative/controller-state.json"
    solo_state = json.loads(solo_state_path.read_text()) if solo_state_path.exists() else None
    state = solo_state or colocated_state
    measured_cells = [
        x for x in cells.values() if x["measurement_status"] == "SOLO_AUTHORITATIVE"
    ]
    coverage = {
        "solo_primary_anchors_at_least_25": len(all_primary) >= 25,
        "solo_measured_cells_12": len(measured_cells) == 12,
    }
    invariants = {
        "surface_rows_12": len(cells) == 12,
        "selection_counts_match": all(selection_matches),
        "client_invariants": all(all_client_invariants),
        "cell_validity": all(
            all(x["cell_valid"]["invariants"].values()) for x in measured_cells
        ),
        "gpu_below_6": float(state["gpu_hours_total"]) < 6.0,
        "rates_nonnegative": all(x["f24"] is None or x["f24"] >= 0 for x in cells.values()),
        "surface_not_identical": len(set(scores24.values())) > 1,
    }
    red_flags = [key for key, value in {**coverage, **invariants}.items() if not value]
    drifted = [key for key, value in cells.items() if value["material_frontier_moved"] is True]
    flip_cells = [
        key for key, value in cells.items()
        if any(anchor["feasibility_flip"] for anchor in value["anchors"])
    ]
    partial = bool([key for key, value in coverage.items() if not value])
    w1_audit_path = root / "w1-readjudication-A-P6-1.json"
    w1_audit = json.loads(w1_audit_path.read_text()) if w1_audit_path.exists() else None
    censored = {key: value["censor"] for key, value in cells.items() if not value["bounded"]}
    colocated_deltas = [
        {
            "cell": cell, "anchor": anchor["anchor"],
            "solo_pass_rate": anchor["accepted_pass_rate"],
            "solo_feasible": anchor["accepted_feasible"],
            **anchor["colocated"],
        }
        for cell, value in cells.items() for anchor in value["anchors"]
        if anchor.get("colocated") is not None
    ]
    return {
        "schema": 1,
        "status": "BUDGET_STOP_PARTIAL" if partial else ("VALID" if not red_flags else "INVALID"),
        "authoritative_tier": "A-P6-2 solo host placement",
        "limitation": "Upgrade-path churn includes dash1->dash0 and resolved-default changes. Co-located W1-W3 values are indicative only; P3 composition is contextual, not a matched v0.20 Layer-1 baseline.",
        "ground_truth_sha256": sha256_file(ground_path), "cells": cells,
        "floor_buckets": {"v20_tol": tol20, "v20": buckets20, "v24_tol": tol24, "v24": buckets24},
        "verdicts": {
            "argmax": argmax, "ranking": ranking, "trap": trap,
            "full_surface_bounded": full_bounded, "tau_b": tau,
            "top_pair_reversals_gt5pct": reversals,
        },
        "materially_drifted_cells": drifted,
        "feasibility_flip_cells": flip_cells,
        "decision_blockers": {
            "coverage": [key for key, value in coverage.items() if not value],
            "unbounded_or_unresolved_cells": censored,
        },
        "solo_vs_colocated": colocated_deltas,
        "w1_readjudication": w1_audit,
        "run_stats": {
            "measured_cells": len(measured_cells),
            "surface_cells": len(cells),
            "primary_anchor_runs": len(all_primary),
            "confirmation_runs": sum(
                len(anchor["confirmations"])
                for cell in cells.values() for anchor in cell["anchors"]
            ),
            "accepted_anchor_trials": sum(
                anchor["trial_count"]
                for cell in cells.values() for anchor in cell["anchors"]
            ),
            "warmup_runs": len(measured_cells),
            "solo_cell_gpu_hours": {
                key: value.get("gpu_hours") for key, value in (solo_state or {}).get("cells", {}).items()
            },
            "colocated_wave_gpu_hours": {
                key: value.get("gpu_hours") for key, value in colocated_state["waves"].items()
            },
            "launch_echo": (
                (root / "launch-echo.log").read_text().splitlines()
                + ((root / "solo-authoritative/launch-echo.log").read_text().splitlines()
                   if (root / "solo-authoritative/launch-echo.log").exists() else [])
            ),
        },
        "attempt_history": {
            "colocated_status": colocated_state["status"],
            "colocated_h20_hours": colocated_state["gpu_hours_total"],
            "colocated_primary_anchors": colocated_state["completed_primary_anchors"],
            "colocated_confirmations": colocated_state["completed_confirmations"],
            "colocated_budget_stop": colocated_state.get("budget_stop"),
            "solo_status": (solo_state or {}).get("status"),
            "solo_repairs": (solo_state or {}).get("repairs", []),
            "solo_failures": (solo_state or {}).get("failures", []),
            "raw_roots": {
                "colocated": str(root / "cells"),
                "solo_authoritative": str(root / "solo-authoritative/cells"),
            },
        },
        "p3_composition_reference": p3_reference(p3_root),
        "gpu": {
            "new_h20_hours": state["gpu_hours_total"], "hard_cap": 6.0,
            "prior_colocated_h20_hours": colocated_state["gpu_hours_total"],
            "solo_h20_hours": (solo_state or {}).get("solo_gpu_hours", 0.0),
            "completed_primary_anchors": (solo_state or {}).get("primary_anchors", 0),
            "confirmations": (solo_state or {}).get("confirmations", 0),
            "controller_status": state["status"],
            "budget_stop": state.get("budget_stop"),
        },
        "sanity": {
            "red_flags": red_flags, "coverage": coverage, "invariants": invariants,
            "numeric": {
                "v20_score": numeric(list(scores20.values())),
                "v24_score": numeric(list(scores24.values())),
                "drift": numeric([x["drift"] for x in cells.values()]),
                "primary_pass_rate": numeric([x["pass_rate"] for x in all_primary]),
                "selected_count": numeric([x["selection"]["count"] for x in all_primary]),
                "layer1_steps": numeric([
                    anchor["layer1"]["steps"] for cell in cells.values() for anchor in cell["anchors"]
                ]),
            },
        },
    }


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--root", type=Path, required=True)
    p.add_argument("--ground-truth", type=Path, required=True)
    p.add_argument("--p3-root", type=Path, required=True)
    p.add_argument("--out", type=Path, required=True)
    args = p.parse_args()
    result = analyze(args.root, args.ground_truth, args.p3_root)
    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"], "verdicts": result["verdicts"], "red_flags": result["sanity"]["red_flags"], "gpu": result["gpu"]}, sort_keys=True))


if __name__ == "__main__":
    main()
