#!/usr/bin/env python3
"""Toy verification of saturation-adaptive prefill chunking.

This is intentionally a mechanism test, not a GPU-serving benchmark.  Each
quantum has a saturation/whale workload state and a prefill pulse.  Chunk size
controls the pulse ramp; a bounded demand envelope keeps peak power mostly
independent of chunking, as predicted by the proposal.
"""
import json
import math
import numpy as np

SEED = 1663
EPS = 0.05
DT = 1.0
C_MIN, C_MAX = 8.0, 64.0
TARGET = 20.0


def controller(s, w, ks=28.0, kw=20.0):
    """c_t = clip(cmax - ks*s - kw*w, cmin, cmax)."""
    return np.clip(C_MAX - ks*np.asarray(s) - kw*np.asarray(w), C_MIN, C_MAX)


def trace(s, w, chunks, seed=SEED):
    """Generate power samples and return ramp and peak reserve statistics.

    The high-frequency noise is deliberately small.  The prefill pulse is
    proportional to chunk size and workload pressure, while the demand
    envelope (and hence the dominant peak) depends on workload rather than
    chunk size.
    """
    rng = np.random.default_rng(seed)
    s, w, chunks = map(np.asarray, (s, w, chunks))
    pressure = 0.35 + 0.65*s + 0.80*w
    # A workload envelope gives similar peak power for different chunking.
    envelope = 500.0 + 35.0*s + 9.0*w
    # c/64 is an incremental prefill ramp, not a new peak envelope.
    pulse = 17.0 * (chunks/C_MAX) * pressure
    noise = rng.normal(0.0, 0.35, size=len(s))
    power = envelope + pulse + noise
    ramps = np.abs(np.diff(power)) / DT
    return power, ramps


def q(x, p=1-EPS):
    return float(np.quantile(x, p, method="linear"))


def workload(n, s_level, whale_fraction, seed=SEED):
    rng = np.random.default_rng(seed + int(1000*s_level) + int(100*whale_fraction))
    # Correlated but nonconstant load makes the ramp statistic meaningful.
    s = np.clip(s_level + rng.normal(0, 0.055, n), 0, 1)
    w = (rng.random(n) < whale_fraction).astype(float)
    # whale requests also slightly increase instantaneous saturation
    s = np.clip(s + 0.04*w, 0, 1)
    return s, w


def run():
    results = {"seed": SEED, "epsilon": EPS, "target": TARGET}

    # Math/controller sanity: exact slope before clipping, and monotonicity.
    s_grid = np.linspace(0, 1, 101)
    c0 = controller(s_grid, np.zeros_like(s_grid))
    c_slope = float(np.polyfit(s_grid[(s_grid > .1)&(s_grid < .8)], c0[(s_grid > .1)&(s_grid < .8)], 1)[0])
    w_grid = np.linspace(0, 1, 101)
    c_w = controller(np.full_like(w_grid, .2), w_grid)
    results["controller_check"] = {
        "interior_slope_observed": c_slope,
        "interior_slope_predicted": -28.0,
        "s_monotone": bool(np.all(np.diff(c0) <= 1e-12)),
        "w_monotone": bool(np.all(np.diff(c_w) <= 1e-12)),
        "clipped_range": [float(c0.min()), float(c0.max())],
    }

    # Prediction 1: controller response is affine until the lower bound.
    n = 12000
    s, w = workload(n, .55, .30)
    fixed = np.full(n, C_MAX)
    adaptive = controller(s, w)
    p_fixed, r_fixed = trace(s, w, fixed, seed=11)
    p_adapt, r_adapt = trace(s, w, adaptive, seed=11)
    fixed_q95 = q(r_fixed)
    adaptive_q95 = q(r_adapt)
    fixed_peak = float(np.quantile(p_fixed, .99))
    adaptive_peak = float(np.quantile(p_adapt, .99))
    results["main_comparison"] = {
        "fixed_chunk": C_MAX,
        "adaptive_mean_chunk": float(adaptive.mean()),
        "adaptive_min_max_chunk": [float(adaptive.min()), float(adaptive.max())],
        "fixed_ramp_q95": fixed_q95,
        "adaptive_ramp_q95": adaptive_q95,
        "ramp_reduction_percent": 100*(1-adaptive_q95/fixed_q95),
        "fixed_peak_q99": fixed_peak,
        "adaptive_peak_q99": adaptive_peak,
        "peak_change_percent": 100*(adaptive_peak/fixed_peak-1),
        "throughput_proxy_fixed_tokens_per_quantum": C_MAX,
        "throughput_proxy_adaptive_tokens_per_quantum": float(adaptive.mean()),
        "throughput_proxy_change_percent": 100*(adaptive.mean()/C_MAX-1),
        "latency_proxy_quantum_ratio_adaptive_over_fixed": float(C_MAX/adaptive.mean()),
    }

    # Prediction 2: with fixed workload state, ramp reserve rises with chunk.
    sweep = []
    s2, w2 = workload(10000, .75, .50, seed=77)
    for c in [8, 16, 32, 48, 64]:
        power, ramps = trace(s2, w2, np.full(len(s2), c), seed=22)
        sweep.append({"chunk": c, "ramp_q95": q(ramps), "peak_q99": float(np.quantile(power,.99))})
    ramp_slope = float(np.polyfit([x["chunk"] for x in sweep], [x["ramp_q95"] for x in sweep], 1)[0])
    results["chunk_sweep"] = {"points": sweep, "q95_ramp_slope_per_token": ramp_slope,
                               "predicted_sign": "positive"}

    # Prediction 3: adaptive benefit grows with saturation and whale load.
    regimes = []
    for sat in [.15, .50, .85]:
        for whale in [0.0, .5]:
            ss, ww = workload(9000, sat, whale, seed=300 + int(100*sat)+int(whale*10))
            ff = np.full(len(ss), C_MAX)
            aa = controller(ss, ww)
            _, rf = trace(ss, ww, ff, seed=31)
            _, ra = trace(ss, ww, aa, seed=31)
            regimes.append({"saturation": sat, "whale_fraction": whale,
                            "adaptive_mean_chunk": float(aa.mean()),
                            "fixed_ramp_q95": q(rf), "adaptive_ramp_q95": q(ra),
                            "reduction_percent": 100*(1-q(ra)/q(rf))})
    results["regime_sweep"] = regimes
    high = [x for x in regimes if x["saturation"] == .85 and x["whale_fraction"] == .5][0]
    low = [x for x in regimes if x["saturation"] == .15 and x["whale_fraction"] == 0.0][0]

    # Quantile selection rule: largest tested c under target reserve.
    candidates = [x for x in sweep if x["ramp_q95"] <= TARGET]
    results["quantile_selection"] = {"target_ramp": TARGET,
        "largest_feasible_chunk": max((x["chunk"] for x in candidates), default=None),
        "candidate_reserves": {str(x["chunk"]): x["ramp_q95"] for x in sweep}}

    # Quantitative predictions stated explicitly for auditability.
    results["predictions_observed_vs_predicted"] = {
        "P1_controller_slope_vs_formula": {"predicted": -28.0, "observed": c_slope, "units": "tokens per saturation unit"},
        "P2_ramp_reserve_chunk_slope_positive": {"predicted": "> 0", "observed": ramp_slope, "units": "q95 ramp units per token"},
        "P3_high_load_reduction_exceeds_low_load": {"predicted": "high > low", "observed_high_percent": high["reduction_percent"], "observed_low_percent": low["reduction_percent"]},
        "P4_peak_change_is_small": {"predicted": "absolute change < 5%", "observed_percent": results["main_comparison"]["peak_change_percent"]}
    }

    # Explicit checks used for the final decision.
    results["checks"] = {
        "controller_formula_pass": abs(c_slope + 28) < .15 and results["controller_check"]["s_monotone"],
        "ramp_increases_with_chunk": ramp_slope > 0,
        "high_load_benefit_exceeds_low_load": high["reduction_percent"] > low["reduction_percent"],
        "peak_change_small_main_case": abs(results["main_comparison"]["peak_change_percent"]) < 5.0,
    }
    results["all_mechanism_checks_pass"] = all(results["checks"].values())
    print(json.dumps(results, indent=2, default=lambda x: x.item() if hasattr(x, "item") else str(x)))

if __name__ == "__main__":
    run()