#!/usr/bin/env python3
"""Compare captured vLLM expert loads with Frontier's fixed routing prior."""

from __future__ import annotations

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

import numpy as np


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--routing", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args()


def gini(values: np.ndarray) -> float:
    values = np.asarray(values, dtype=np.float64)
    if values.sum() == 0:
        return 0.0
    ordered = np.sort(values)
    n = len(ordered)
    indices = np.arange(1, n + 1, dtype=np.float64)
    return float((2 * np.sum(indices * ordered) / np.sum(ordered) - (n + 1)) / n)


def stats(values: np.ndarray) -> dict[str, float]:
    values = np.asarray(values, dtype=np.float64)
    mean = float(np.mean(values))
    return {
        "load_cv": float(np.std(values) / mean),
        "load_gini": gini(values),
        "max_load_ratio": float(np.max(values) / mean),
        "expert_utilization": float(np.count_nonzero(values) / len(values)),
    }


def proportional_counts(total: int, ratios: np.ndarray) -> np.ndarray:
    exact = total * ratios / ratios.sum()
    counts = np.floor(exact).astype(np.int64)
    remainder = total - int(counts.sum())
    order = sorted(
        range(len(ratios)),
        key=lambda i: (-(exact[i] - counts[i]), -(ratios[i] / ratios.sum()), i),
    )
    for index in range(remainder):
        counts[order[index % len(order)]] += 1
    assert int(counts.sum()) == total
    return counts


def correlation(left: np.ndarray, right: np.ndarray) -> float:
    value = float(np.corrcoef(left, right)[0, 1])
    return value if math.isfinite(value) else 0.0


def distribution(values: list[float]) -> dict[str, float]:
    return {
        "min": min(values),
        "median": statistics.median(values),
        "max": max(values),
    }


def main() -> None:
    args = parse_args()
    routing = json.loads(args.routing.read_text())
    phases = routing["phases"]
    layer_count = len(phases["prefill"]["per_layer"])
    expert_count = len(phases["prefill"]["per_layer"][0]["counts"])

    rows: list[dict[str, Any]] = []
    for phase in ("prefill", "decode"):
        for layer, actual in enumerate(phases[phase]["per_layer"]):
            rng = np.random.RandomState(args.seed + layer)
            ratios = rng.uniform(0.1, 1.0, expert_count)
            synthetic = proportional_counts(int(actual["total_routed_tokens"]), ratios)
            synthetic_stats = stats(synthetic)
            for name in synthetic_stats:
                if not math.isclose(
                    stats(np.asarray(actual["counts"]))[name],
                    float(actual[name]),
                    rel_tol=0,
                    abs_tol=1e-12,
                ):
                    raise ValueError(f"captured metric mismatch: {phase} layer {layer} {name}")
            rows.append(
                {
                    "phase": phase,
                    "layer": layer,
                    "total_routed_tokens": int(actual["total_routed_tokens"]),
                    "actual": {name: float(actual[name]) for name in synthetic_stats},
                    "frontier_simulation": synthetic_stats,
                    "actual_vs_frontier_pearson": correlation(
                        np.asarray(actual["counts"], dtype=np.float64), synthetic
                    ),
                }
            )

    phase_summary: dict[str, Any] = {}
    for phase in ("prefill", "decode"):
        selected = [row for row in rows if row["phase"] == phase]
        phase_summary[phase] = {
            "token_count": int(phases[phase]["token_count"]),
            "actual": {
                name: distribution([row["actual"][name] for row in selected])
                for name in selected[0]["actual"]
            },
            "frontier_simulation": {
                name: distribution([row["frontier_simulation"][name] for row in selected])
                for name in selected[0]["frontier_simulation"]
            },
            "actual_vs_frontier_pearson": distribution(
                [row["actual_vs_frontier_pearson"] for row in selected]
            ),
        }

    phase_correlations = []
    for layer in range(layer_count):
        prefill = np.asarray(phases["prefill"]["per_layer"][layer]["counts"])
        decode = np.asarray(phases["decode"]["per_layer"][layer]["counts"])
        phase_correlations.append(correlation(prefill, decode))

    output = {
        "schema": "frontier-routing-mismatch.v1",
        "source": str(args.routing.resolve()),
        "frontier_contract": {
            "mode": "simulation",
            "seed": args.seed,
            "allocation": "per-layer fixed Uniform(0.1, 1.0), normalized once and reused for every batch and phase",
            "layer_count": layer_count,
            "expert_count": expert_count,
        },
        "phase_summary": phase_summary,
        "actual_prefill_vs_decode_pearson": distribution(phase_correlations),
        "frontier_prefill_vs_decode_pearson": 1.0,
        "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()