#!/usr/bin/env python3
"""Analyze Qwen235 Fixed-PD real-state proxies and Frontier component ledgers."""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from statistics import fmean
from typing import Any, Iterable


HERE = Path(__file__).resolve().parent
REPO_ROOT = HERE.parents[1]
sys.path.insert(0, str(REPO_ROOT / "runs/telemetry-residual"))

from common_state import load_jsonl, numeric  # noqa: E402


CONFIGS = ("tp4_ep1_mns64", "tp8_ep8_mns64")
TP_BY_CONFIG = {"tp4_ep1_mns64": 4, "tp8_ep8_mns64": 8}
REAL_TPOT_MS = {"tp4_ep1_mns64": 21.0442, "tp8_ep8_mns64": 27.9942}
LOG_PATTERN = re.compile(
    r"Avg prompt throughput: (?P<prompt>[0-9.]+) tokens/s, "
    r"Avg generation throughput: (?P<generation>[0-9.]+) tokens/s, "
    r"Running: (?P<running>[0-9]+) reqs, Waiting: (?P<waiting>[0-9]+) reqs, "
    r"GPU KV cache usage: (?P<kv>[0-9.]+)%"
)
ITERATION_PATTERN = re.compile(
    r"Iteration\((?P<index>[0-9]+)\): (?P<context_requests>[0-9]+) context requests, "
    r"(?P<context_tokens>[0-9]+) context tokens, "
    r"(?P<generation_requests>[0-9]+) generation requests, "
    r"(?P<generation_tokens>[0-9]+) generation tokens, iteration elapsed time: "
    r"(?P<elapsed_ms>[0-9.]+) ms"
)


CATEGORIES = {
    "attention": {
        "attention_prefill_execution_time",
        "attention_decode_execution_time",
        "attention_pre_proj_time",
        "attention_post_proj_time",
        "attention_kv_cache_save_execution_time",
        "attention_rope_execution_time",
        "attn_norm_time",
    },
    "dense_mlp_compute": {
        "mlp_layer_up_proj_execution_time",
        "mlp_layer_act_execution_time",
        "mlp_layer_down_proj_execution_time",
        "mlp_norm_time",
    },
    "moe_compute": {
        "moe_grouped_gemm_time",
        "share_expert_up_proj_time",
        "share_expert_act_time",
        "share_expert_down_proj_time",
    },
    "moe_routing": {
        "moe_gating_linear_time",
        "moe_gating_routing_topk_time",
        "moe_shuffling_time",
    },
    "ep_communication": {"expert_parallel_communication_time"},
    "tp_dp_communication": {
        "attention_all_reduce_time",
        "mlp_all_reduce_time",
        "moe_tensor_parallel_allgather_time",
        "share_expert_tensor_parallel_allreduce_time",
        "dp_input_allreduce_time",
        "dp_output_allreduce_time",
    },
    "pipeline_communication": {"pipeline_parallel_communication_time"},
    "runtime_overhead": {
        "add_attn_residual_time",
        "add_ffn_residual_time",
        "schedule_time",
        "sampler_e2e_time",
        "prepare_inputs_e2e_time",
        "pp_producer_send_path_runtime_time",
        "pp_receiver_head_runtime_time",
        "pp_prefill_consumer_active_runtime_time",
        "pp_stage_boundary_residual_runtime_time",
        "process_model_outputs_time",
        "ray_comm_time",
        "decode_draft_proposer_time",
        "mtp_terminal_overshoot_time",
    },
}

OP_CATEGORIES = {
    "attention": {
        "input_layernorm",
        "attn_pre_proj",
        "attn_rope",
        "attn_prefill",
        "attn_decode",
        "attn_kv_cache_save",
        "attn_post_proj",
    },
    "dense_mlp_compute": {
        "post_attention_layernorm",
        "mlp_up_proj",
        "mlp_act",
        "mlp_down_proj",
    },
    "moe_compute": {
        "moe_grouped_gemm",
        "share_expert_up_proj",
        "share_expert_act",
        "share_expert_down_proj",
    },
    "moe_routing": {
        "moe_gating_linear",
        "moe_gating_routing_topk",
        "moe_shuffling",
    },
    "ep_communication": {
        "expert_parallel_alltoall",
        "expert_parallel_alltoall_dispatch",
        "expert_parallel_alltoall_combine",
        "expert_parallel_allreduce",
    },
    "tp_dp_communication": {
        "attn_tensor_parallel_allreduce",
        "mlp_tensor_parallel_allreduce",
        "moe_tensor_parallel_allgather",
        "moe_tensor_parallel_allreduce",
        "share_expert_tensor_parallel_allreduce",
        "dp_input_allreduce",
        "dp_output_allreduce",
    },
    "pipeline_communication": {"pipeline_parallel_send_recv"},
    "runtime_overhead": {
        "add_attn_residual",
        "add_ffn_residual",
        "schedule",
        "prepare_inputs_e2e",
        "pp_receiver_head_runtime",
        "pp_prefill_consumer_active_runtime",
        "decode_draft_proposer",
        "mtp_terminal_overshoot",
        "pp_stage_boundary_handoff",
        "sampler_e2e",
        "process_model_outputs",
        "ray_comm_time",
    },
}
DIAGNOSTIC_OPS = {"expert_parallel_allreduce_wait"}


def atomic_write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(text)
    os.replace(temporary, path)


def parse_real_log(path: Path) -> list[dict[str, float | int]]:
    rows = []
    for line in path.read_text(errors="replace").splitlines():
        match = LOG_PATTERN.search(line)
        if match is None:
            continue
        rows.append(
            {
                "prompt_tokens_per_s": float(match.group("prompt")),
                "generation_tokens_per_s": float(match.group("generation")),
                "running": int(match.group("running")),
                "waiting": int(match.group("waiting")),
                "kv_percent": float(match.group("kv")),
            }
        )
    if not rows:
        raise ValueError(f"no vLLM periodic metrics found: {path}")
    return rows


def summarize_real(real_root: Path) -> dict[str, Any]:
    result = {}
    for config in CONFIGS:
        tp = TP_BY_CONFIG[config]
        expected_prompt = 4096 * 0.2 * tp
        trials = []
        steady_all = []
        for trial in ("trial1", "trial2", "trial3"):
            path = real_root / config / trial / "logs/server.log"
            rows = parse_real_log(path)
            steady = [
                row
                for row in rows
                if row["prompt_tokens_per_s"] >= 0.95 * expected_prompt
                and row["generation_tokens_per_s"] > 0
            ]
            if len(steady) < 5:
                raise ValueError(f"insufficient steady real proxy samples: {path}")
            steady_all.extend(steady)
            trials.append(
                {
                    "trial": trial,
                    "path": str(path.resolve()),
                    "all_samples": len(rows),
                    "steady_samples": len(steady),
                    "running": numeric(row["running"] for row in steady),
                    "waiting": numeric(row["waiting"] for row in steady),
                    "kv_percent": numeric(row["kv_percent"] for row in steady),
                    "generation_tokens_per_s": numeric(
                        row["generation_tokens_per_s"] for row in steady
                    ),
                }
            )
        running_counts = Counter(int(row["running"]) for row in steady_all)
        result[config] = {
            "proxy_only": True,
            "steady_rule": f"prompt throughput >= 95% of {expected_prompt:.1f} tokens/s",
            "trials": trials,
            "aggregate": {
                "samples": len(steady_all),
                "running": numeric(row["running"] for row in steady_all),
                "running_histogram": {
                    str(key): value for key, value in sorted(running_counts.items())
                },
                "waiting": numeric(row["waiting"] for row in steady_all),
                "kv_percent": numeric(row["kv_percent"] for row in steady_all),
                "generation_tokens_per_s": numeric(
                    row["generation_tokens_per_s"] for row in steady_all
                ),
            },
        }
    return result


def parse_iteration_log(path: Path) -> list[dict[str, float | int]]:
    rows = []
    for line in path.read_text(errors="replace").splitlines():
        match = ITERATION_PATTERN.search(line)
        if match is None:
            continue
        rows.append(
            {
                "index": int(match.group("index")),
                "context_requests": int(match.group("context_requests")),
                "context_tokens": int(match.group("context_tokens")),
                "generation_requests": int(match.group("generation_requests")),
                "generation_tokens": int(match.group("generation_tokens")),
                "elapsed_ms": float(match.group("elapsed_ms")),
            }
        )
    if not rows:
        raise ValueError(f"no iteration details found: {path}")
    start = next(
        (index for index, row in enumerate(rows) if row["context_tokens"] >= 4096),
        None,
    )
    if start is None:
        raise ValueError(f"measured Fixed-PD interval not found: {path}")
    measured = rows[start:]
    if any(row["generation_tokens"] != row["generation_requests"] for row in measured):
        raise ValueError("Fixed-PD decode must schedule one token per generation request")
    return measured


def summarize_iteration_real(iteration_root: Path) -> dict[str, Any]:
    result = {}
    for config in CONFIGS:
        path = iteration_root / config / "logs/server.log"
        rows = parse_iteration_log(path)
        decode = [row for row in rows if row["generation_requests"] > 0]
        pure = [row for row in decode if row["context_tokens"] == 0]
        mixed = [row for row in decode if row["context_tokens"] > 0]
        if not pure or not mixed:
            raise ValueError(f"iteration state lacks pure or mixed decode rows: {path}")
        joint = Counter(
            (int(row["context_tokens"]), int(row["generation_requests"]))
            for row in decode
        )
        token_joint = {
            f"{context}:{generation}": count * generation
            for (context, generation), count in sorted(joint.items())
        }
        pure_hist = Counter(int(row["generation_requests"]) for row in pure)
        weights = [int(row["generation_requests"]) for row in decode]
        result[config] = {
            "path": str(path.resolve()),
            "measured_rows": len(rows),
            "decode_bearing_rows": len(decode),
            "pure_decode_rows": len(pure),
            "mixed_prefill_decode_rows": len(mixed),
            "decode_batch_size": numeric(
                row["generation_requests"] for row in decode
            ),
            "pure_decode_batch_histogram": {
                str(key): value for key, value in sorted(pure_hist.items())
            },
            "decode_token_weighted_joint_state_histogram": token_joint,
            "decode_token_weighted_iteration_elapsed_ms": sum(
                float(row["elapsed_ms"]) * weight
                for row, weight in zip(decode, weights, strict=True)
            )
            / sum(weights),
            "pure_decode_iteration_elapsed_ms": numeric(
                row["elapsed_ms"] for row in pure
            ),
            "mixed_iteration_elapsed_ms": numeric(
                row["elapsed_ms"] for row in mixed
            ),
        }
    return result


def categorized_components(components: dict[str, Any]) -> dict[str, float]:
    covered = set().union(*CATEGORIES.values())
    unknown = set(components) - covered
    missing = covered - set(components)
    if unknown or missing:
        raise ValueError(
            f"component schema drift: unknown={sorted(unknown)}, missing={sorted(missing)}"
        )
    return {
        category: sum(float(components[name]) for name in names)
        for category, names in CATEGORIES.items()
    }


def op_category(name: str) -> str | None:
    if name in DIAGNOSTIC_OPS:
        return None
    matches = [category for category, names in OP_CATEGORIES.items() if name in names]
    if len(matches) != 1:
        raise ValueError(f"unclassified or multiply classified op trace event: {name}")
    return matches[0]


def load_stage_rows(state_root: Path, config: str) -> tuple[list[dict[str, Any]], dict]:
    result_path = state_root / config / "result.json"
    result = json.loads(result_path.read_text())
    if result.get("status") != "PASS":
        raise ValueError(f"state replay did not pass: {result_path}")
    if not result["scorer_equivalence"]["request_metrics_byte_identical"]:
        raise ValueError(f"state replay scorer changed: {result_path}")
    ledger = Path(result["state_artifacts"]["ledger"]["path"])
    rows = []
    for row in load_jsonl(ledger):
        token_counts = [int(value) for value in row["request_num_tokens"]]
        if not token_counts:
            raise ValueError(f"empty stage batch in {config}")
        components = categorized_components(
            row["execution_time"]["component_ledger_ms"]
        )
        total = float(row["execution_time"]["total_time_ms"])
        if not math.isclose(sum(components.values()), total, abs_tol=1e-6):
            raise ValueError(f"categorized components do not sum for {config}")
        rows.append(
            {
                "batch_size": len(row["request_ids"]),
                "decode_requests": sum(value == 1 for value in token_counts),
                "prefill_requests": sum(value > 1 for value in token_counts),
                "prefill_tokens": sum(value for value in token_counts if value > 1),
                "total_time_ms": total,
                "categories_ms": components,
                "per_expert_tokens": row.get("per_expert_tokens"),
            }
        )
    if not rows:
        raise ValueError(f"no decode-only ledger rows for {config}")
    return rows, result


def load_op_trace_rows(
    op_trace_root: Path, config: str
) -> tuple[list[dict[str, Any]], dict, dict[str, Any]]:
    result_path = op_trace_root / config / "result.json"
    result = json.loads(result_path.read_text())
    if result.get("status") != "PASS" or not result.get("op_trace_enabled"):
        raise ValueError(f"op-trace replay did not pass: {result_path}")
    trace_path = Path(result["state_artifacts"]["op_trace"]["path"])
    grouped: dict[int, dict[str, Any]] = {}
    event_names = Counter()
    with trace_path.open() as source:
        for line_number, line in enumerate(source, start=1):
            event = json.loads(line)
            if "meta" in event and len(event) == 1:
                continue
            name = str(event["name"])
            category = op_category(name)
            event_names[name] += 1
            batch_id = int(event["batch_id"])
            tokens = [int(value) for value in event["meta"]["num_tokens"]]
            entry = grouped.setdefault(
                batch_id,
                {
                    "tokens": tokens,
                    "categories_ms": {key: 0.0 for key in CATEGORIES},
                    "diagnostic_wait_ms": 0.0,
                    "events": 0,
                },
            )
            if entry["tokens"] != tokens:
                raise ValueError(
                    f"op trace token metadata changed within batch {batch_id}"
                )
            if category is None:
                entry["diagnostic_wait_ms"] += float(event["duration_ms"])
            else:
                entry["categories_ms"][category] += float(event["duration_ms"])
            entry["events"] += 1
    if not grouped:
        raise ValueError(f"no op trace events: {trace_path}")

    ledger_path = Path(result["state_artifacts"]["ledger"]["path"])
    stage_span_by_batch = {
        int(row["batch_id"]):
        (float(row["stage_end_ts"]) - float(row["stage_start_ts"])) * 1000
        for row in load_jsonl(ledger_path)
    }
    rows = []
    for batch_id, entry in sorted(grouped.items()):
        tokens = entry["tokens"]
        if batch_id not in stage_span_by_batch:
            raise ValueError(f"op trace batch missing from ledger: {batch_id}")
        rows.append(
            {
                "batch_id": batch_id,
                "batch_size": len(tokens),
                "decode_requests": sum(value == 1 for value in tokens),
                "prefill_requests": sum(value > 1 for value in tokens),
                "prefill_tokens": sum(value for value in tokens if value > 1),
                "total_time_ms": stage_span_by_batch[batch_id],
                "categories_ms": entry["categories_ms"],
                "diagnostic_wait_ms": entry["diagnostic_wait_ms"],
                "per_expert_tokens": None,
            }
        )

    ratios = [
        sum(row["categories_ms"].values())
        / stage_span_by_batch[row["batch_id"]]
        for row in rows
        if row["batch_id"] in stage_span_by_batch
        and stage_span_by_batch[row["batch_id"]] > 0
    ]
    validation = {
        "trace_path": str(trace_path),
        "trace_batches": len(rows),
        "ledger_batches": len(stage_span_by_batch),
        "serialized_component_sum_over_stage_span": numeric(ratios),
        "diagnostic_wait_ms": numeric(
            row["diagnostic_wait_ms"]
            for row in rows
            if row["diagnostic_wait_ms"] > 0
        ),
        "event_names": dict(sorted(event_names.items())),
    }
    return rows, result, validation


def weighted_component_mean(
    rows: Iterable[dict[str, Any]], *, weight_name: str = "batch_size"
) -> dict[str, float]:
    selected = list(rows)
    if not selected:
        raise ValueError("component mean needs rows")
    weights = [int(row[weight_name]) for row in selected]
    if any(weight <= 0 for weight in weights):
        raise ValueError(f"component weights must be positive: {weight_name}")
    denominator = sum(weights)
    values = {
        category: sum(
            row["categories_ms"][category] * weight
            for row, weight in zip(selected, weights, strict=True)
        )
        / denominator
        for category in CATEGORIES
    }
    values["total"] = sum(
        row["total_time_ms"] * weight
        for row, weight in zip(selected, weights, strict=True)
    ) / denominator
    return values


def means_by_batch(rows: list[dict[str, Any]]) -> dict[int, dict[str, Any]]:
    groups: dict[int, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        groups[row["batch_size"]].append(row)
    return {
        batch_size: {
            "n": len(group),
            "components_ms": weighted_component_mean(group),
        }
        for batch_size, group in sorted(groups.items())
    }


def means_by_joint_state(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        if row["decode_requests"] <= 0:
            continue
        key = f"{row['prefill_tokens']}:{row['decode_requests']}"
        groups[key].append(row)
    return {
        key: {
            "n": len(group),
            "components_ms": weighted_component_mean(
                group, weight_name="decode_requests"
            ),
        }
        for key, group in sorted(groups.items())
    }


def reweight(
    by_batch: dict[int, dict[str, Any]], histogram: dict[str, int]
) -> dict[str, Any]:
    total_samples = sum(histogram.values())
    supported = {
        int(batch_size): count
        for batch_size, count in histogram.items()
        if int(batch_size) in by_batch
    }
    covered = sum(supported.values())
    if covered == 0:
        return {"coverage": 0.0, "components_ms": None, "unsupported": histogram}
    values = {
        category: sum(
            by_batch[batch_size]["components_ms"][category] * count
            for batch_size, count in supported.items()
        )
        / covered
        for category in (*CATEGORIES, "total")
    }
    return {
        "coverage": covered / total_samples,
        "supported_samples": covered,
        "total_samples": total_samples,
        "components_ms": values,
        "unsupported": {
            batch_size: count
            for batch_size, count in histogram.items()
            if int(batch_size) not in by_batch
        },
    }


def reweight_joint(
    by_state: dict[str, dict[str, Any]], token_histogram: dict[str, int]
) -> dict[str, Any]:
    total_tokens = sum(token_histogram.values())
    supported = {
        state: count
        for state, count in token_histogram.items()
        if state in by_state
    }
    covered = sum(supported.values())
    if covered == 0:
        return {
            "coverage": 0.0,
            "components_ms": None,
            "unsupported": token_histogram,
        }
    return {
        "coverage": covered / total_tokens,
        "supported_decode_tokens": covered,
        "total_decode_tokens": total_tokens,
        "components_ms": {
            category: sum(
                by_state[state]["components_ms"][category] * count
                for state, count in supported.items()
            )
            / covered
            for category in (*CATEGORIES, "total")
        },
        "unsupported": {
            state: count
            for state, count in token_histogram.items()
            if state not in by_state
        },
    }
def subtract(right: dict[str, float], left: dict[str, float]) -> dict[str, float]:
    return {name: right[name] - left[name] for name in left}


def analyze(
    real_root: Path,
    state_root: Path,
    op_trace_root: Path | None = None,
    iteration_root: Path | None = None,
) -> dict[str, Any]:
    real = summarize_real(real_root)
    iteration_real = (
        summarize_iteration_real(iteration_root) if iteration_root is not None else None
    )
    sim = {}
    by_batch = {}
    by_joint_state = {}
    for config in CONFIGS:
        component_source = "stage_batch_ledger"
        trace_validation = None
        if op_trace_root is not None and (op_trace_root / config / "result.json").is_file():
            rows, replay, trace_validation = load_op_trace_rows(op_trace_root, config)
            component_source = "op_trace_execution_time_override"
        else:
            rows, replay = load_stage_rows(state_root, config)
        decode_bearing = [row for row in rows if row["decode_requests"] > 0]
        pure_decode = [row for row in decode_bearing if row["prefill_requests"] == 0]
        mixed = [row for row in decode_bearing if row["prefill_requests"] > 0]
        if not pure_decode:
            raise ValueError(f"no pure-decode ledger rows for {config}")
        grouped = means_by_batch(pure_decode)
        by_batch[config] = grouped
        by_joint_state[config] = means_by_joint_state(decode_bearing)
        sim[config] = {
            "stage_rows": len(rows),
            "decode_bearing_rows": len(decode_bearing),
            "pure_decode_rows": len(pure_decode),
            "mixed_prefill_decode_rows": len(mixed),
            "decode_batch_size": numeric(
                row["decode_requests"] for row in decode_bearing
            ),
            "decode_token_weighted_all_step_components_ms": weighted_component_mean(
                decode_bearing, weight_name="decode_requests"
            ),
            "decode_token_weighted_mixed_step_share": sum(
                row["decode_requests"] for row in mixed
            )
            / sum(row["decode_requests"] for row in decode_bearing),
            "pure_decode_token_weighted_components_ms": weighted_component_mean(
                pure_decode, weight_name="decode_requests"
            ),
            "batch_support": {
                str(batch): value for batch, value in grouped.items()
            },
            "scorer_equivalence": replay["scorer_equivalence"],
            "component_source": component_source,
            "trace_validation": trace_validation,
        }

    real_reweighted = {
        config: reweight(
            by_batch[config], real[config]["aggregate"]["running_histogram"]
        )
        for config in CONFIGS
    }
    exact_reweighted = None
    exact_contrast = None
    if iteration_real is not None:
        exact_reweighted = {
            config: reweight_joint(
                by_joint_state[config],
                iteration_real[config][
                    "decode_token_weighted_joint_state_histogram"
                ],
            )
            for config in CONFIGS
        }
        if min(value["coverage"] for value in exact_reweighted.values()) >= 0.8:
            exact_contrast = subtract(
                exact_reweighted[CONFIGS[1]]["components_ms"],
                exact_reweighted[CONFIGS[0]]["components_ms"],
            )
    simulator_internal_contrast = subtract(
        sim[CONFIGS[1]]["decode_token_weighted_all_step_components_ms"],
        sim[CONFIGS[0]]["decode_token_weighted_all_step_components_ms"],
    )
    coverages = [real_reweighted[config]["coverage"] for config in CONFIGS]
    proxy_contrast = None
    if min(coverages) >= 0.8:
        proxy_contrast = subtract(
            real_reweighted[CONFIGS[1]]["components_ms"],
            real_reweighted[CONFIGS[0]]["components_ms"],
        )

    shared = sorted(set(by_batch[CONFIGS[0]]) & set(by_batch[CONFIGS[1]]))
    shared_contrasts = {
        str(batch): {
            "tp4_n": by_batch[CONFIGS[0]][batch]["n"],
            "tp8_n": by_batch[CONFIGS[1]][batch]["n"],
            "tp8_minus_tp4_ms": subtract(
                by_batch[CONFIGS[1]][batch]["components_ms"],
                by_batch[CONFIGS[0]][batch]["components_ms"],
            ),
        }
        for batch in shared
        if min(by_batch[config][batch]["n"] for config in CONFIGS) >= 10
    }
    observed_real_contrast = REAL_TPOT_MS[CONFIGS[1]] - REAL_TPOT_MS[CONFIGS[0]]
    exact_iteration_contrast = None
    if iteration_real is not None:
        exact_iteration_contrast = (
            iteration_real[CONFIGS[1]][
                "decode_token_weighted_iteration_elapsed_ms"
            ]
            - iteration_real[CONFIGS[0]][
                "decode_token_weighted_iteration_elapsed_ms"
            ]
        )
    decision_contrast = exact_contrast if exact_contrast is not None else proxy_contrast
    if decision_contrast is None:
        verdict = "STOP: real Running proxy has insufficient exact simulator support"
    elif decision_contrast["total"] < 0:
        if exact_contrast is not None:
            verdict = (
                "State-composition mismatch is insufficient: after reweighting Frontier "
                "to measured real decode composition, it still predicts TP8 faster."
            )
        else:
            verdict = (
                "Active-batch-count mismatch is insufficient: after reweighting Frontier "
                "to the frozen-real Running proxy, it still predicts TP8 faster."
            )
    else:
        if exact_contrast is not None:
            verdict = (
                "Exact real state composition flips Frontier to the correct TP4 winner; "
                "the ranking error is caused by closed-loop state/composition mismatch, "
                "not a conditional stage reversal."
            )
        else:
            verdict = (
                "Coarse active-batch state can flip Frontier's ordering, but "
                "iteration-level batch/context telemetry is required before attributing "
                "the real gap to state."
            )
    decomposition = None
    if exact_contrast is not None:
        decomposition = {
            "simulator_internal_tp8_minus_tp4_ms": simulator_internal_contrast["total"],
            "exact_state_matched_tp8_minus_tp4_ms": exact_contrast["total"],
            "state_composition_shift_ms": (
                exact_contrast["total"] - simulator_internal_contrast["total"]
            ),
            "residual_vs_frozen_real_tpot_ms": (
                exact_contrast["total"] - observed_real_contrast
            ),
            "residual_vs_exact_iteration_ms": (
                exact_contrast["total"] - exact_iteration_contrast
            ),
        }
    return {
        "schema": "qwen235-fixed-pd-state-diagnosis-v1",
        "status": "PASS",
        "scope": {
            "workload": "Fixed-PD 4096->256, 0.2 req/s/GPU, MNS64",
            "real_proxy_limitation": (
                "vLLM 10-second Running is active requests, not per-iteration decode batch; "
                "context lengths and graph buckets are unavailable in frozen logs"
            ),
            "graph_observability": (
                "Frontier state outputs have no direct graph bucket/padding/launch-overhead "
                "field; graph effects remain folded into predictors"
            ),
            "op_trace_accounting": (
                "For TP8 shared-domain sync, total is ledger stage_end-start (critical path); "
                "op categories are serialized work estimates and are not additive because "
                "overlap and lane-summed wait diagnostics are represented separately"
            ),
        },
        "real": real,
        "real_iteration_state": iteration_real,
        "simulator": sim,
        "simulator_internal_all_step_tp8_minus_tp4_ms": simulator_internal_contrast,
        "proxy_matched": {
            "method": (
                "exact batch-size lookup; each simulator config is reweighted to its own "
                "frozen-real steady Running histogram; no interpolation"
            ),
            "configs": real_reweighted,
            "tp8_minus_tp4_ms": proxy_contrast,
        },
        "exact_state_matched": {
            "method": (
                "decode-token-weighted exact (context_tokens, generation_requests) "
                "composition from vLLM iteration details; no interpolation"
            ),
            "configs": exact_reweighted,
            "tp8_minus_tp4_ms": exact_contrast,
        }
        if iteration_real is not None
        else None,
        "same_batch_contrasts": shared_contrasts,
        "contrast_decomposition": decomposition,
        "reference": {
            "real_tpot_ms": REAL_TPOT_MS,
            "observed_real_tp8_minus_tp4_ms": observed_real_contrast,
            "exact_iteration_tp8_minus_tp4_ms": exact_iteration_contrast,
        },
        "verdict": verdict,
    }


def markdown(result: dict[str, Any]) -> str:
    lines = [
        "# Qwen235 Fixed-PD state diagnosis",
        "",
        f"**Verdict:** {result['verdict']}",
        "",
        "| Config | Real Running proxy mean | Sim decode batch mean | Proxy coverage |",
        "|---|---:|---:|---:|",
    ]
    for config in CONFIGS:
        lines.append(
            f"| {config} | {result['real'][config]['aggregate']['running']['mean']:.3f} "
            f"| {result['simulator'][config]['decode_batch_size']['mean']:.3f} "
            f"| {result['proxy_matched']['configs'][config]['coverage']:.1%} |"
        )
    exact = result["exact_state_matched"]
    iteration = result["real_iteration_state"]
    if exact is not None:
        lines.extend(
            [
                "",
                "| Config | Real exact decode batch mean | Real token-weighted iteration ms | Matched simulator ms | Exact coverage |",
                "|---|---:|---:|---:|---:|",
            ]
        )
        for config in CONFIGS:
            lines.append(
                f"| {config} | {iteration[config]['decode_batch_size']['mean']:.3f} "
                f"| {iteration[config]['decode_token_weighted_iteration_elapsed_ms']:.4f} "
                f"| {exact['configs'][config]['components_ms']['total']:.4f} "
                f"| {exact['configs'][config]['coverage']:.1%} |"
            )
    contrast = (
        exact["tp8_minus_tp4_ms"]
        if exact is not None and exact["tp8_minus_tp4_ms"] is not None
        else result["proxy_matched"]["tp8_minus_tp4_ms"]
    )
    internal = result["simulator_internal_all_step_tp8_minus_tp4_ms"]
    lines.extend(
        [
            "",
            "## Frontier internal component contrast over its own executed composition",
            "",
            "Decode-token-weighted over both pure-decode and mixed prefill/decode steps.",
            "",
            "| Component | TP8 - TP4 (ms/decoded token step) |",
            "|---|---:|",
        ]
    )
    for name, value in sorted(
        internal.items(), key=lambda item: abs(item[1]), reverse=True
    ):
        lines.append(f"| {name} | {value:+.4f} |")
    if contrast is not None:
        heading = (
            "Frontier component contrast at exact real token composition"
            if exact is not None and exact["tp8_minus_tp4_ms"] is not None
            else "Frontier internal component contrast at real Running proxy"
        )
        lines.extend(
            [
                "",
                f"## {heading}",
                "",
                "Positive means TP8 slower; negative means Frontier gives TP8 an advantage.",
                "",
                "| Component | TP8 - TP4 (ms/step) |",
                "|---|---:|",
            ]
        )
        for name, value in sorted(
            contrast.items(), key=lambda item: abs(item[1]), reverse=True
        ):
            lines.append(f"| {name} | {value:+.4f} |")
    lines.extend(
        [
            "",
            "## Interpretation boundary",
            "",
            f"- Real observed TPOT contrast: {result['reference']['observed_real_tp8_minus_tp4_ms']:+.4f} ms/token.",
        ]
    )
    if result["reference"]["exact_iteration_tp8_minus_tp4_ms"] is not None:
        lines.append(
            "- Real exact-iteration contrast: "
            f"{result['reference']['exact_iteration_tp8_minus_tp4_ms']:+.4f} "
            "ms/decode-token-weighted step."
        )
    lines.extend(
        [
            f"- {result['scope']['real_proxy_limitation']}.",
            f"- {result['scope']['graph_observability']}.",
            f"- {result['scope']['op_trace_accounting']}.",
            "- Component deltas identify where Frontier creates its own TP8 advantage; without real per-stage measurements they are not yet root-cause proof.",
            "",
        ]
    )
    return "\n".join(lines)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--real-root", type=Path, required=True)
    parser.add_argument("--state-root", type=Path, required=True)
    parser.add_argument("--op-trace-root", type=Path)
    parser.add_argument("--iteration-root", type=Path)
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    result = analyze(
        args.real_root.resolve(),
        args.state_root.resolve(),
        args.op_trace_root.resolve() if args.op_trace_root else None,
        args.iteration_root.resolve() if args.iteration_root else None,
    )
    atomic_write(args.json_output, json.dumps(result, indent=2, sort_keys=True) + "\n")
    atomic_write(args.markdown_output, markdown(result))
    print(json.dumps({"status": result["status"], "verdict": result["verdict"]}))


if __name__ == "__main__":
    main()
