#!/usr/bin/env python3
"""Compare Frontier operator predictions at the identical initial scheduler state."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path


PATTERN = re.compile(
    r"\[OP-TRACE\]\[MONOLITHIC\]\[(?:ATTENTION|MOE)\]\[([^]]+)\] "
    r"batch_id=(\d+), layer_id=(\d+), predicted_time_ms=([0-9.eE+-]+)"
)
COMPONENTS = (
    "input_layernorm",
    "attn_pre_proj",
    "attn_rope",
    "attn_kv_cache_save",
    "attn_prefill",
    "attn_decode",
    "attn_post_proj",
    "post_attention_layernorm",
    "moe_gating_linear",
    "moe_gating_routing_topk",
    "moe_shuffling",
    "moe_grouped_gemm",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--old", type=Path, required=True)
    parser.add_argument("--new", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def first_state(path: Path) -> dict[str, float]:
    values: dict[str, set[float]] = {}
    with path.open(errors="replace") as handle:
        for line in handle:
            match = PATTERN.search(line)
            if match and int(match[2]) == 0 and int(match[3]) == 0:
                values.setdefault(match[1], set()).add(float(match[4]))
    ambiguous = {name: sorted(items) for name, items in values.items() if len(items) != 1}
    if ambiguous:
        raise ValueError(f"ambiguous initial predictions in {path}: {ambiguous}")
    result = {name: next(iter(items)) for name, items in values.items()}
    result.setdefault("attn_decode", 0.0)
    missing = sorted(set(COMPONENTS) - set(result))
    if missing:
        raise ValueError(f"missing initial predictions in {path}: {missing}")
    return result


def main() -> None:
    args = parse_args()
    old = first_state(args.old)
    new = first_state(args.new)
    rows = []
    for component in COMPONENTS:
        rows.append(
            {
                "component": component,
                "historical_profile_ms": old[component],
                "vllm020_profile_ms": new[component],
                "new_over_old": new[component] / old[component]
                if old[component] != 0
                else None,
            }
        )
    old_total = sum(old[name] for name in COMPONENTS)
    new_total = sum(new[name] for name in COMPONENTS)
    output = {
        "schema": "frontier-initial-op-trace-delta.v1",
        "comparison_contract": {
            "fixture": "fidelity_p1_tp1_mns64_low1",
            "scheduler_state": "batch_id=0, layer_id=0 before profile-dependent trajectories diverge",
            "calibration_a_tp": 1.0,
        },
        "layer_component_sum_ms": {
            "historical_profile": old_total,
            "vllm020_profile": new_total,
            "new_over_old": new_total / old_total,
        },
        "rows": rows,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
    print(args.output)


if __name__ == "__main__":
    main()