src/inferscale/consolidation.py

from __future__ import annotations

import math
import random
from copy import deepcopy
from statistics import mean, median
from typing import Any

from .execution import ExecutionLearningConfig, generate_workflows, run_execution_learning


POLICIES = (
    ("No prefetch", "none"),
    ("Top-1 decayed", "decayed"),
    ("Multi-step top-k", "multistep"),
    ("Utility-aware multi-step", "utility"),
)


def _percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(float(value) for value in values)
    if len(ordered) == 1:
        return ordered[0]
    pos = min(max(float(q), 0.0), 1.0) * (len(ordered) - 1)
    lo = int(math.floor(pos))
    hi = int(math.ceil(pos))
    if lo == hi:
        return ordered[lo]
    frac = pos - lo
    return ordered[lo] * (1.0 - frac) + ordered[hi] * frac


def _bootstrap_mean_ci(values: list[float], samples: int, seed: int) -> tuple[float, float]:
    if not values:
        return (0.0, 0.0)
    if len(values) == 1:
        return (values[0], values[0])
    rng = random.Random(seed)
    n = len(values)
    draws = []
    for _ in range(max(100, int(samples))):
        draws.append(mean(values[rng.randrange(n)] for _ in range(n)))
    return _percentile(draws, 0.025), _percentile(draws, 0.975)


def _summary_row(label: str, result: dict[str, Any]) -> dict[str, Any]:
    return {
        "label": label,
        "p95_ttft_ms": float(result["latency"]["step_ttft_ms"]["p95"]),
        "p95_workflow_e2e_ms": float(result["latency"]["workflow_e2e_ms"]["p95"]),
        "workflow_throughput_rps": float(result["summary"]["workflow_throughput_rps"]),
        "completion_rate": float(result["summary"]["workflow_completion_rate"]),
        "prefix_hit_rate": float(result["resource"]["prefix_hit_rate"]),
        "prefetch_utilization": float(result["resource"].get("prefetch_utilization", 0.0)),
        "future_role_recall_at_k": float(result["resource"].get("forecast_recall", 0.0)),
        "mean_hbm_gb": float(result["resource"]["mean_prefix_hbm_gb"]),
        "unused_prefetch_gb": float(result["resource"].get("unused_prefetch_gb", 0.0)),
        "pressure_evictions": int(result["resource"].get("pressure_evictions", 0)),
        "saved_prefill_tokens": int(result["resource"].get("prefill_tokens_saved", 0)),
    }


def _dominates(a: dict[str, Any], b: dict[str, Any]) -> bool:
    # The robust frontier deliberately treats speculative traffic and HBM as
    # first-class resources rather than ranking on latency alone.
    a_obj = (float(a["p95_ttft_ms"]), float(a["unused_prefetch_gb"]), float(a["mean_hbm_gb"]))
    b_obj = (float(b["p95_ttft_ms"]), float(b["unused_prefetch_gb"]), float(b["mean_hbm_gb"]))
    return all(x <= y + 1e-12 for x, y in zip(a_obj, b_obj, strict=True)) and any(
        x < y - 1e-12 for x, y in zip(a_obj, b_obj, strict=True)
    )


def _pareto_labels(rows: list[dict[str, Any]]) -> set[str]:
    labels: set[str] = set()
    for candidate in rows:
        if not any(_dominates(other, candidate) for other in rows if other is not candidate):
            labels.add(str(candidate["label"]))
    return labels


def _offline_constrained_oracle(
    base: ExecutionLearningConfig,
    workflows: list[Any],
    deployable_results: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    """Return a bounded full-trace information upper bound.

    This is intentionally *not* described as a globally optimal cache controller.
    It is an exhaustive oracle over a declared candidate family: all deployable
    policies already evaluated plus clairvoyant future-set plans for horizons
    1..5 and top-k 1..3. Every candidate uses the same cache budget, transfer
    bandwidth, model, device profile, and exact realized workflow trace.
    """
    pool: list[dict[str, Any]] = []
    for label, result in deployable_results.items():
        pool.append(
            {
                "label": label,
                "kind": "deployable",
                "config": {
                    "policy": result["config"]["prefetch_policy"],
                    "horizon": result["config"].get("forecast_horizon", base.forecast_horizon),
                    "top_k": result["config"].get("prefetch_top_k", base.prefetch_top_k),
                },
                "result": result,
                "row": _summary_row(label, result),
            }
        )

    for horizon in range(1, 6):
        for top_k in range(1, 4):
            cfg = ExecutionLearningConfig.from_dict(base.to_dict())
            cfg.prefetch_policy = "oracle_horizon"
            cfg.forecast_horizon = horizon
            cfg.prefetch_top_k = top_k
            result = run_execution_learning(cfg.to_dict(), workflows)
            label = f"clairvoyant H{horizon}/K{top_k}"
            pool.append(
                {
                    "label": label,
                    "kind": "clairvoyant",
                    "config": {"policy": "oracle_horizon", "horizon": horizon, "top_k": top_k},
                    "result": result,
                    "row": _summary_row(label, result),
                }
            )

    feasible = [item for item in pool if item["row"]["completion_rate"] >= 1.0 - 1e-12]
    if not feasible:
        feasible = pool
    winner = min(
        feasible,
        key=lambda item: (
            item["row"]["p95_ttft_ms"],
            item["row"]["p95_workflow_e2e_ms"],
            item["row"]["unused_prefetch_gb"],
            item["row"]["mean_hbm_gb"],
        ),
    )
    return {
        "label": winner["label"],
        "kind": winner["kind"],
        "config": winner["config"],
        "metrics": winner["row"],
        "candidate_count": len(pool),
        "definition": "bounded-full-trace-serving-oracle",
        "note": (
            "Exhaustive upper bound over the declared candidate family, including clairvoyant future-set plans. "
            "It uses the complete trace for policy selection and future-role actions, but is not a proof of global optimality."
        ),
    }


def repeated_seed_policy_study(
    config: dict[str, Any],
    repetitions: int = 12,
    bootstrap_samples: int = 600,
) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    repetitions = max(4, min(int(repetitions), 24))
    bootstrap_samples = max(100, min(int(bootstrap_samples), 4000))
    seed_runs: list[dict[str, Any]] = []
    per_policy: dict[str, list[dict[str, Any]]] = {label: [] for label, _ in POLICIES}

    for rep in range(repetitions):
        seed = base.seed + rep * 1009
        seeded = ExecutionLearningConfig.from_dict(base.to_dict())
        seeded.seed = seed
        workflows = generate_workflows(seeded)
        deployable_results: dict[str, dict[str, Any]] = {}
        rows: list[dict[str, Any]] = []
        for label, policy in POLICIES:
            cfg = ExecutionLearningConfig.from_dict(seeded.to_dict())
            cfg.prefetch_policy = policy
            result = run_execution_learning(cfg.to_dict(), workflows)
            deployable_results[label] = result
            row = _summary_row(label, result)
            rows.append(row)
            per_policy[label].append(row)

        pareto = _pareto_labels(rows)
        nominal_winner = min(rows, key=lambda row: (row["p95_ttft_ms"], row["unused_prefetch_gb"]))["label"]
        oracle = _offline_constrained_oracle(seeded, workflows, deployable_results)
        oracle_ttft = float(oracle["metrics"]["p95_ttft_ms"])
        for row in rows:
            row["pareto"] = row["label"] in pareto
            row["ttft_winner"] = row["label"] == nominal_winner
            row["oracle_regret_ms"] = max(0.0, float(row["p95_ttft_ms"]) - oracle_ttft)
            row["oracle_regret_pct"] = (
                row["oracle_regret_ms"] / oracle_ttft * 100.0 if oracle_ttft > 1e-12 else 0.0
            )
        seed_runs.append({"rep": rep + 1, "seed": seed, "rows": rows, "oracle": oracle})

    baseline_label = "Top-1 decayed"
    baseline_ttfts = [float(row["p95_ttft_ms"]) for row in per_policy[baseline_label]]
    summaries: list[dict[str, Any]] = []
    for index, (label, _) in enumerate(POLICIES):
        rows = per_policy[label]
        ttfts = [float(row["p95_ttft_ms"]) for row in rows]
        ci_low, ci_high = _bootstrap_mean_ci(ttfts, bootstrap_samples, base.seed ^ (index + 1) * 7919)
        paired_deltas = [ttft - base_ttft for ttft, base_ttft in zip(ttfts, baseline_ttfts, strict=True)]
        delta_low, delta_high = _bootstrap_mean_ci(
            paired_deltas, bootstrap_samples, base.seed ^ (index + 1) * 104729
        )
        seed_rows = [
            next(row for row in seed_run["rows"] if row["label"] == label)
            for seed_run in seed_runs
        ]
        regrets = [float(row["oracle_regret_ms"]) for row in seed_rows]
        regret_low, regret_high = _bootstrap_mean_ci(
            regrets, bootstrap_samples, base.seed ^ (index + 1) * 15485863
        )
        summaries.append(
            {
                "label": label,
                "mean_ttft_ms": mean(ttfts),
                "median_ttft_ms": median(ttfts),
                "ttft_ci95_low_ms": ci_low,
                "ttft_ci95_high_ms": ci_high,
                "paired_delta_vs_top1_mean_ms": mean(paired_deltas),
                "paired_delta_ci95_low_ms": delta_low,
                "paired_delta_ci95_high_ms": delta_high,
                "ttft_win_rate": mean(1.0 if row["ttft_winner"] else 0.0 for row in seed_rows),
                "pareto_stability": mean(1.0 if row["pareto"] else 0.0 for row in seed_rows),
                "median_oracle_regret_ms": median(regrets),
                "mean_oracle_regret_ms": mean(regrets),
                "oracle_regret_ci95_low_ms": regret_low,
                "oracle_regret_ci95_high_ms": regret_high,
                "worst_seed_ttft_ms": max(ttfts),
                "mean_unused_prefetch_gb": mean(float(row["unused_prefetch_gb"]) for row in rows),
                "mean_hbm_gb": mean(float(row["mean_hbm_gb"]) for row in rows),
                "mean_completion_rate": mean(float(row["completion_rate"]) for row in rows),
            }
        )

    ranked = sorted(
        summaries,
        key=lambda row: (
            -float(row["ttft_win_rate"]),
            float(row["median_ttft_ms"]),
            float(row["median_oracle_regret_ms"]),
            -float(row["pareto_stability"]),
            float(row["mean_unused_prefetch_gb"]),
        ),
    )
    rank_map = {row["label"]: rank + 1 for rank, row in enumerate(ranked)}
    for row in summaries:
        row["robust_rank"] = rank_map[row["label"]]

    nominal_first_seed = min(
        seed_runs[0]["rows"], key=lambda row: (row["p95_ttft_ms"], row["unused_prefetch_gb"])
    )["label"]
    oracle_ttfts = [float(seed_run["oracle"]["metrics"]["p95_ttft_ms"]) for seed_run in seed_runs]
    oracle_labels: dict[str, int] = {}
    for seed_run in seed_runs:
        label = str(seed_run["oracle"]["label"])
        oracle_labels[label] = oracle_labels.get(label, 0) + 1

    return {
        "study": "repeated-seed-policy-consolidation",
        "protocol": "matched-seeds-bootstrap-and-bounded-offline-oracle",
        "repetitions": repetitions,
        "bootstrap_samples": bootstrap_samples,
        "config": base.to_dict(),
        "policies": summaries,
        "seed_runs": seed_runs,
        "nominal_winner_first_seed": nominal_first_seed,
        "robust_winner": ranked[0]["label"] if ranked else None,
        "oracle": {
            "definition": "bounded-full-trace-serving-oracle",
            "candidate_count_per_seed": seed_runs[0]["oracle"]["candidate_count"] if seed_runs else 0,
            "median_ttft_ms": median(oracle_ttfts) if oracle_ttfts else 0.0,
            "mean_ttft_ms": mean(oracle_ttfts) if oracle_ttfts else 0.0,
            "selected_plan_frequency": oracle_labels,
            "note": (
                "The oracle exhaustively selects among the evaluated deployable policies plus clairvoyant future-set "
                "plans over H=1..5 and K=1..3 on each complete trace, under the same cache/bandwidth constraints. "
                "It is a bounded information upper bound, not a proof of globally optimal action scheduling."
            ),
        },
        "pareto_objectives": ["p95_step_ttft_ms", "unused_prefetch_gb", "mean_hbm_gb"],
        "note": (
            "Robust rank prioritizes how often a policy wins TTFT across matched seeds, then median TTFT and regret to "
            "the bounded offline oracle. Pareto stability is reported separately rather than silently overriding latency. "
            "Bootstrap intervals quantify seed uncertainty, not real-hardware error."
        ),
    }