"""Grid-flexibility (demand-response) analysis for LLM serving fleets.

Models the GPU-to-Grid (G2G) batch-size control mechanism proposed by
Hassan et al. (arXiv:2602.05116v1, 2025): when the grid signals high load,
the LLM serving engine caps the maximum in-flight batch size (vLLM
``max_num_seqs``).  A lower cap reduces GPU power draw at the cost of
higher queuing latency.

The main entry point :func:`grid_flex_analysis` sweeps power-reduction
percentages and returns a list of :class:`GridFlexPoint` trade-off entries,
each carrying analytical and optional DES-verified P99 TTFT estimates.
For each flex level it: (1) inverts the GPU power model to derive
``n_max_cap``; (2) estimates P99 TTFT analytically (M/G/c); (3) optionally
runs a DES simulation to verify the analytical estimate.

Throttled recalibration
-----------------------
At n_max_cap < n_slots_full, each GPU iteration covers fewer concurrent
sequences: iter_t(n_max_cap) = W + H*n_max_cap < iter_t(n_slots_full).
Using the full-load mu_slot underestimates throughput and makes the
analytical model falsely diverge (inf P99) when n_max_cap is below
the full-load Erlang saturation threshold but actually still stable.
Using a throttled profile gives the correct M/G/c calibration for the
reduced concurrency level.

Verification
------------
Both analytical and DES paths predict P99 TTFT independently; agreement
(typically within 10–15%) confirms the M/G/c model is adequate for
production sizing.  The *power* side is verified by comparing the logistic
power curve against the ML.ENERGY benchmark data points embedded in the
profile constants.  The G2G paper uses ``H100-SXM5`` hardware; the
``H100_80GB`` profile logistic parameters (k=1.0, x0=4.2) are fitted to
those measurements.

Sources
-------
- Hassan et al. (2025) "GPU-to-Grid: Coupling LLM Inference with Power
  System Control", arXiv:2602.05116v1.
- ML.ENERGY Benchmark v3.0 (Chung et al., NeurIPS 2025 D&B track),
  https://ml.energy/leaderboard — measured H100-SXM5 power vs. batch size.
- NVIDIA GPU Spec Sheets: H100-SXM5 TDP 700W; A100-SXM4 TDP 400W.

Examples
--------
>>> from fleet_sim import grid_flex_analysis, H100_80GB
>>> import json
>>> cdf = json.load(open("data/azure_cdf.json"))
>>> # Analytical only (fast)
>>> curve = grid_flex_analysis(cdf, lam=200, n_gpus=40, gpu=H100_80GB,
...                            t_slo_ms=500)
>>> # With DES verification (slower but verified)
>>> curve = grid_flex_analysis(cdf, lam=200, n_gpus=40, gpu=H100_80GB,
...                            t_slo_ms=500, n_sim_requests=20000)
>>> for pt in curve:
...     des = f"{pt.p99_ttft_des_ms:.1f}ms" if pt.p99_ttft_des_ms else "—"
...     print(f"flex={pt.flex_pct:4.0f}%  n_max={pt.n_max_cap:3d}"
...           f"  {pt.power_per_gpu_w:.0f}W/GPU  P99={pt.p99_ttft_ms:.1f}ms"
...           f"  DES={des}  {'OK' if pt.slo_met else 'BREACH'}")
"""

import dataclasses
import math
from dataclasses import dataclass

from ..core.fleet import Fleet, FleetConfig, PoolConfig
from ..gpu_profiles.manual import ManualProfile
from ..workload.synthetic import PoissonWorkload
from . import analytical

# ── Grid-flexibility analysis ─────────────────────────────────────────────────


@dataclass
class GridFlexPoint:
    """One operating point on the power-vs-latency trade-off curve.

    Generated by :func:`grid_flex_analysis`.  Each row answers:
    "If I cap per-GPU batch size to *n_max_cap* concurrent requests
    (committing *flex_pct*% power reduction to the grid), what P99 TTFT
    does the fleet deliver?"

    Attributes
    ----------
    flex_pct          : power reduction committed to the grid operator (%)
    n_max_cap         : resulting concurrent-request cap per GPU
    power_per_gpu_w   : estimated per-GPU power draw at this cap (W)
    power_fleet_kw    : total fleet power (n_gpus × power_per_gpu_w / 1000)
    p99_ttft_ms       : analytical P99 TTFT estimate at this cap (ms)
    p99_ttft_des_ms   : DES-verified P99 TTFT (ms); None if DES not run
    slo_met           : True if p99 ≤ t_slo_ms (uses DES P99 when available)
    power_model       : "logistic" or "linear" — which GPU power model was used
    """

    flex_pct: float
    n_max_cap: int
    power_per_gpu_w: float
    power_fleet_kw: float
    p99_ttft_ms: float
    slo_met: bool
    p99_ttft_des_ms: float | None = None
    power_model: str = "linear"


def _throttled_profile(gpu, n_max_cap: int, max_ctx: int):
    """Return a copy of *gpu* with KV-slot count forced to *n_max_cap*.

    Used by grid_flex_analysis() DES verification to simulate a fleet with
    a batch-size cap applied (the G2G max_num_seqs curtailment mechanism).
    """

    blks_per_seq = math.ceil(max_ctx / gpu.blk_size)
    # Force both KV-memory limit and compute limit to equal n_max_cap
    new_kv = n_max_cap * blks_per_seq
    # compute_cap = max_slots * calibration_ctx // max_ctx; we want this = n_max_cap
    new_max_slots = math.ceil(n_max_cap * max_ctx / max(1, gpu.calibration_ctx))
    return dataclasses.replace(gpu, max_slots=new_max_slots, total_kv_blks=new_kv)


def grid_flex_analysis(
    cdf: list,
    lam: float,
    n_gpus: int,
    gpu: ManualProfile,
    t_slo_ms: float,
    max_ctx: int = 8192,
    flex_pcts: list[float] | None = None,
    n_sim_requests: int = 0,
    verbose: bool = False,
) -> list[GridFlexPoint]:
    """Compute the power–latency trade-off curve for a fleet under demand response.

    Parameters
    ----------
    cdf             : workload CDF as list of (token_threshold, cumulative_frac)
    lam             : total arrival rate (req/s) hitting this fleet
    n_gpus          : fixed fleet size to evaluate (e.g. from a prior optimize() run)
    gpu             : ManualProfile with power_idle_w and power_nominal_w set
    t_slo_ms        : P99 TTFT SLO (ms) — used to flag slo_met
    max_ctx         : maximum context window for this pool (tokens)
    flex_pcts       : power-reduction pcts to sweep, default [0, 5, ..., 50].
    n_sim_requests  : DES request count per flex level. 0 (default) disables DES.
                      10k–30k for stable P99. Independent per level, scales linearly.
    verbose         : print progress when running DES verification.

    Returns
    -------
    List of :class:`GridFlexPoint`, one per flex_pct, sorted ascending.
    Each point includes ``p99_ttft_des_ms`` when DES is enabled.

    Raises
    ------
    ValueError – if *gpu* has no power model constants configured.
    """
    if flex_pcts is None:
        flex_pcts = [0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0]

    # Validate power model is available
    _ = gpu.power_at_concurrency(1)  # raises ValueError if not configured
    power_model_type = "logistic" if gpu.power_logistic_k > 0.0 else "linear"
    n_slots_full = gpu.n_slots(max_ctx)
    results: list[GridFlexPoint] = []
    for flex_pct in flex_pcts:
        # Target power per GPU after flex commitment (clamped to idle floor)
        target_power = max(
            gpu.power_idle_w, gpu.power_nominal_w * (1.0 - flex_pct / 100.0)
        )
        # Invert power model to find n_max_cap, cap at KV-cache slot limit.
        n_max_cap = min(n_slots_full, gpu.invert_power(target_power))
        actual_power = gpu.power_at_concurrency(n_max_cap)
        fleet_kw = n_gpus * actual_power / 1000.0
        # Recalibrate at throttled concurrency (see module docstring).
        capped_gpu = _throttled_profile(gpu, n_max_cap, max_ctx)
        mu_cap, cv2_cap, _, mean_prefill_cap = analytical.calibrate(
            cdf, max_ctx, capped_gpu
        )
        mu_slot_cap = mu_cap / n_max_cap if n_max_cap > 0 else 1e-9
        c_slots = n_gpus * n_max_cap
        p99_s = analytical.p99_wait(c_slots, lam, mu_slot_cap, cv2_cap)
        p99_ttft_ms = (p99_s + mean_prefill_cap) * 1000.0

        # ── Optional DES verification ─────────────────────────────────────
        p99_ttft_des_ms = None
        if n_sim_requests > 0:
            try:
                des_p99 = _run_grid_flex_des(
                    cdf=cdf,
                    lam=lam,
                    n_gpus=n_gpus,
                    gpu=gpu,
                    n_max_cap=n_max_cap,
                    max_ctx=max_ctx,
                    n_req=n_sim_requests,
                )
                p99_ttft_des_ms = des_p99
                if verbose:
                    print(
                        f"  flex={flex_pct:.0f}%  n_max={n_max_cap}"
                        f"  analytical={p99_ttft_ms:.1f}ms"
                        f"  DES={p99_ttft_des_ms:.1f}ms"
                    )
            except Exception as e:
                if verbose:
                    print(f"  [DES error at flex={flex_pct}%]: {e}")
        # slo_met uses DES P99 when available (more accurate)
        effective_p99 = p99_ttft_des_ms if p99_ttft_des_ms is not None else p99_ttft_ms
        results.append(
            GridFlexPoint(
                flex_pct=flex_pct,
                n_max_cap=n_max_cap,
                power_per_gpu_w=actual_power,
                power_fleet_kw=fleet_kw,
                p99_ttft_ms=p99_ttft_ms,
                p99_ttft_des_ms=p99_ttft_des_ms,
                slo_met=effective_p99 <= t_slo_ms,
                power_model=power_model_type,
            )
        )
    return results


def _run_grid_flex_des(
    cdf: list,
    lam: float,
    n_gpus: int,
    gpu,
    n_max_cap: int,
    max_ctx: int,
    n_req: int,
    seed: int = 42,
) -> float:
    """Run a single-pool DES with ``n_max_cap`` KV slots per GPU.

    Returns P99 TTFT in milliseconds.  Raises if the queue is severely
    overloaded (>95% of requests never complete).
    """
    from ..workload.synthetic import CdfWorkload

    # Build a profile with n_max_cap slots (simulates max_num_seqs cap)
    capped = _throttled_profile(gpu, n_max_cap, max_ctx)

    pool = PoolConfig(pool_id="flex", gpu=capped, n_gpus=n_gpus, max_ctx=max_ctx)
    fc = FleetConfig(pools=[pool], router_type="LengthRouter")

    wl = CdfWorkload(cdf, seed=seed)
    arrivals = PoissonWorkload(lam, wl, n_requests=n_req, seed=seed).generate()

    fleet = Fleet(fc)
    result = fleet.run(arrivals, verbose=False)

    # Fleet has a single pool; use the top-level P99 helper
    p99 = result.p99_ttft_ms()
    return p99 if p99 != float("inf") else float("inf")


def print_grid_flex_table(
    results: list[GridFlexPoint],
    t_slo_ms: float,
    n_gpus: int,
    lam: float,
) -> None:
    """Pretty-print the grid flex analysis results."""
    baseline = results[0] if results else None
    has_des = any(pt.p99_ttft_des_ms is not None for pt in results)
    power_model = results[0].power_model if results else "linear"

    print(f"\n{'='*70}")
    print(f"  Grid Flexibility Analysis  [{power_model} power model]")
    print(f"  Fleet: {n_gpus} GPUs  λ={lam:.0f} req/s  SLO={t_slo_ms:.0f} ms")
    if baseline:
        print(
            f"  Baseline: {baseline.power_fleet_kw:.1f} kW fleet power"
            f"  ({baseline.power_per_gpu_w:.0f} W/GPU)"
        )
    print(f"{'='*70}")

    if has_des:
        print(
            f"  {'Flex':>6} {'n_max':>6} {'W/GPU':>7} {'Fleet kW':>9}"
            f" {'P99 analyt':>11} {'P99 DES':>9} {'SLO':>5}"
        )
        print(f"  {'-'*6} {'-'*6} {'-'*7} {'-'*9} {'-'*11} {'-'*9} {'-'*5}")
        for pt in results:
            ok = "  OK" if pt.slo_met else "BREACH"
            des_str = (
                f"{pt.p99_ttft_des_ms:>8.1f}ms"
                if pt.p99_ttft_des_ms is not None
                else "        —"
            )
            print(
                f"  {pt.flex_pct:>5.0f}% {pt.n_max_cap:>6d}"
                f" {pt.power_per_gpu_w:>6.0f}W {pt.power_fleet_kw:>8.1f}kW"
                f" {pt.p99_ttft_ms:>10.1f}ms {des_str} {ok}"
            )
    else:
        print(
            f"  {'Flex':>6} {'n_max':>6} {'W/GPU':>7} {'Fleet kW':>9}"
            f" {'P99 TTFT':>9} {'SLO':>5}"
        )
        print(f"  {'-'*6} {'-'*6} {'-'*7} {'-'*9} {'-'*9} {'-'*5}")
        for pt in results:
            ok = "  OK" if pt.slo_met else "BREACH"
            print(
                f"  {pt.flex_pct:>5.0f}% {pt.n_max_cap:>6d}"
                f" {pt.power_per_gpu_w:>6.0f}W {pt.power_fleet_kw:>8.1f}kW"
                f" {pt.p99_ttft_ms:>8.1f}ms {ok}"
            )

    # Find max safe flex depth
    safe = [pt for pt in results if pt.slo_met]
    if safe:
        max_safe = max(pt.flex_pct for pt in safe)
        max_safe_pt = max(safe, key=lambda p: p.flex_pct)
        baseline_kw = baseline.power_fleet_kw if baseline else 0.0
        saved_kw = baseline_kw - max_safe_pt.power_fleet_kw
        p99_disp = (
            f"{max_safe_pt.p99_ttft_des_ms:.1f}ms (DES)"
            if max_safe_pt.p99_ttft_des_ms is not None
            else f"{max_safe_pt.p99_ttft_ms:.1f}ms (analytical)"
        )
        print(
            f"\n  Max safe flex depth: {max_safe:.0f}%"
            f"  (saves {saved_kw:.1f} kW fleet-wide,"
            f" P99={p99_disp})"
        )
    else:
        print("\n  No flex depth meets SLO — consider adding GPUs first.")
