"""Tokens-per-watt (energy-efficiency) analysis for GPU serving fleets.

Computes output tok/J across GPU profiles and multi-pool routing topologies.

**Model dependency warning**: W, H, and the power logistic curve all
depend on the served model, not just the GPU.  Pre-built profiles
(H100_80GB, A100_80GB, A10G) are calibrated for Llama-3-70B on 8-GPU TP
(H100/A100) or 7B-class on single-GPU (A10G).  Do not compare ``tok/W``
across profiles that represent different models — the ratio will reflect
model size, not GPU efficiency.

**Single-pool N_gpus cancellation**: For a *single homogeneous pool*,
``tokens_per_watt = mu_gpu × ρ × mean_L_out / P(n_active)`` is
independent of fleet size (N numerator and denominator cancel).

**Multi-pool fleets**: N does NOT cancel across pools.  Use
:class:`FleetTpwResult` which computes the correct aggregate:
``fleet_tpw = Σ(λ_i × L_out_i) / Σ(N_i × P_i(n_i))``.

:func:`fleet_tpw_analysis` is the correct function for comparing:

* **Homo pool**   — all traffic to one GPU type / one model.
* **Hetero pool** — short requests routed to small-GPU/small-model pool,
                    long requests to large-GPU/large-model pool.
* **Semantic router** — fraction α of requests to a fast/cheap pool
                        (e.g. A10G + 7B model), rest to a capable pool
                        (e.g. H100 + 70B).

The CDF for each pool must be normalised so its last cumulative fraction
is 1.0.  Use :func:`_split_cdf` to produce correctly normalised sub-CDFs.
"""

import math
from dataclasses import dataclass

from . import analytical

# ── Tokens-per-Watt analysis ──────────────────────────────────────────────────


@dataclass
class TpwPoint:
    """Energy-efficiency snapshot for one pool at one utilisation level.

    Generated by :func:`tpw_analysis` (single-pool) or per-pool entries
    inside :class:`FleetTpwResult` (multi-pool).

    **Model dependency warning**: W, H, and the power logistic curve all
    depend on the served model, not just the GPU.  Pre-built profiles
    (H100_80GB, A100_80GB, A10G) are calibrated for Llama-3-70B on 8-GPU TP
    (H100/A100) or 7B-class on single-GPU (A10G).  Do not compare ``tok/W``
    across profiles that represent different models — the ratio will reflect
    model size, not GPU efficiency.

    **Single-pool N_gpus cancellation**: For a *single homogeneous pool*,
    ``tokens_per_watt = mu_gpu × ρ × mean_L_out / P(n_active)`` is
    independent of fleet size (N numerator and denominator cancel).

    **Multi-pool fleets**: N does NOT cancel across pools.  Use
    :class:`FleetTpwResult` which computes the correct aggregate:
    ``fleet_tpw = Σ(λ_i × L_out_i) / Σ(N_i × P_i(n_i))``.

    Attributes
    ----------
    gpu_name          : profile label (e.g. "H100-80GB")
    pool_label        : optional human label for the pool (e.g. "short-7B")
    n_gpus            : fleet size at this operating point
    rho               : GPU-level utilisation  λ / (N × mu_gpu)
    n_active          : mean concurrent in-flight requests per GPU
    power_per_gpu_w   : per-GPU power (W) at n_active; from logistic model
    tokens_per_watt   : output tok/J for THIS pool in isolation
                        = (λ_pool × mean_L_out) / (N_gpus × power_per_gpu_w)
    cost_per_mil      : $/1M output tokens for this pool's λ_pool at N_gpus
    p99_ttft_ms       : analytical P99 TTFT (ms) for this pool
    power_model_qual  : "HIGH" / "FAIR" / "LOW" — confidence in P(n_active)
    slo_optimal       : True when N_gpus is the SLO-optimal (minimum) fleet
    """

    gpu_name: str
    n_gpus: int
    rho: float
    n_active: float
    power_per_gpu_w: float
    tokens_per_watt: float
    cost_per_mil: float
    p99_ttft_ms: float
    power_model_qual: str
    pool_label: str = ""
    slo_optimal: bool = False


@dataclass
class FleetTpwResult:
    """Fleet-level tokens-per-watt result for a multi-pool routing topology.

    Generated by :func:`fleet_tpw_analysis`.  Holds per-pool breakdowns and
    the correct fleet-level aggregate.

    Fleet-level formula (N does NOT cancel across pools):

        fleet_tpw = Σ_i(λ_i × mean_L_out_i) / Σ_i(N_i × P_i(n_i))

    This is the only correct formula for comparing homo-pool vs two-pool
    routing, or for semantic-router configurations where different pools
    serve different model sizes.

    Attributes
    ----------
    topology          : human-readable description (e.g. "homo H100-70B",
                        "two-pool A10G-7B + H100-70B")
    pools             : per-pool TpwPoint list (one entry per pool)
    fleet_tpw         : fleet-level tok/W using the formula above
    fleet_cost_per_mil: fleet-level $/1M output tokens
    fleet_power_kw    : total fleet power draw (kW) at this operating point
    total_gpus        : total GPU count across all pools
    worst_p99_ms      : max(p99_ttft per pool)
    power_model_notes : list of caveat strings for LOW-quality power models
    """

    topology: str
    pools: list[TpwPoint]
    fleet_tpw: float
    fleet_cost_per_mil: float
    fleet_power_kw: float
    total_gpus: int
    worst_p99_ms: float
    power_model_notes: list[str]


def _power_model_quality(gpu) -> str:
    """Return 'HIGH' / 'FAIR' / 'LOW' from the profile's own docstring labels."""
    if gpu.power_logistic_k <= 0.0 or gpu.power_idle_w <= 0.0:
        return "NONE"
    name = getattr(gpu, "name", "").upper()
    if "H100" in name or "H200" in name:
        return "HIGH"
    if "A100" in name:
        return "FAIR"
    return "LOW"


def _cdf_mean_l_out(cdf: list) -> int:
    """Return mean output-token count from a (possibly sub-)CDF.

    Assumes the 80/20 input-to-output split used throughout the simulator.
    Works on a normalised sub-CDF (last cumulative fraction == 1.0) or on
    the raw full CDF (last fraction may equal anything ≤ 1.0).
    """
    total_prob = 0.0
    mean_total = 0.0
    prev_t, prev_f = 0, 0.0
    for thresh, frac in cdf:
        width = thresh - prev_t
        prob = frac - prev_f
        mean_total += (prev_t + width / 2.0) * prob
        total_prob += prob
        prev_t, prev_f = thresh, frac
    if total_prob > 0:
        mean_total /= total_prob
    return max(1, int(0.20 * mean_total))


def _split_cdf(full_cdf: list, b_short: int) -> tuple:
    """Split a full CDF into (short_sub_cdf, long_sub_cdf, alpha).

    Parameters
    ----------
    full_cdf : list of (token_threshold, cumulative_frac).
    b_short  : token-count threshold.  Requests with total_tokens ≤ b_short
               go to the short pool; the rest go to the long pool.

    Returns
    -------
    short_cdf : normalised sub-CDF for the short pool (fractions rescaled so
                the last entry is 1.0).
    long_cdf  : normalised sub-CDF for the long pool.
    alpha     : fraction of traffic assigned to the short pool.

    Notes
    -----
    If b_short falls between two CDF knots, linear interpolation is used to
    find alpha accurately.
    """
    # Linear interpolation to find alpha = CDF(b_short)
    alpha = 0.0
    for i, (t, f) in enumerate(full_cdf):
        if t >= b_short:
            if i == 0:
                alpha = f * (b_short / max(1, t))
            else:
                t0, f0 = full_cdf[i - 1]
                if t == t0:
                    alpha = f
                else:
                    alpha = f0 + (f - f0) * (b_short - t0) / (t - t0)
            break
    else:
        alpha = full_cdf[-1][1]  # b_short > max threshold

    alpha = max(1e-6, min(1.0 - 1e-6, alpha))

    short_cdf: list = []
    long_cdf: list = []
    for thresh, frac in full_cdf:
        if thresh <= b_short:
            short_cdf.append((thresh, frac / alpha))
        else:
            norm_f = max(0.0, (frac - alpha) / (1.0 - alpha))
            long_cdf.append((thresh, norm_f))

    # Ensure the short CDF endpoint is exactly b_short at frac=1.0
    if not short_cdf or short_cdf[-1][0] < b_short:
        short_cdf.append((b_short, 1.0))
    else:
        short_cdf[-1] = (short_cdf[-1][0], 1.0)

    if long_cdf:
        long_cdf[-1] = (long_cdf[-1][0], 1.0)
    else:
        # Degenerate: all traffic is short
        long_cdf = [(b_short + 1, 1.0)]

    return short_cdf, long_cdf, alpha


def _tpw_one_pool(
    cdf: list,
    lam: float,
    gpu,
    t_slo_ms: float,
    max_ctx: int,
    pool_label: str = "",
    rho_override: float | None = None,
) -> TpwPoint:
    """Compute a single TpwPoint for one pool.

    Parameters
    ----------
    cdf           : (possibly sub-)CDF for this pool's traffic.
    lam           : arrival rate into THIS pool (req/s).
    gpu           : GPU profile for this pool.
    t_slo_ms      : P99 TTFT SLO (ms).
    max_ctx       : context window (tokens).
    pool_label    : human label (e.g. "short-7B").
    rho_override  : if set, force this utilisation instead of SLO-optimal.
    """
    mu_gpu, cv2, n_slots, mean_pref = analytical.calibrate(cdf, max_ctx, gpu)
    t_slo_s = t_slo_ms / 1000.0
    qual = _power_model_quality(gpu)
    mean_l_out = _cdf_mean_l_out(cdf)

    if rho_override is None:
        n_gpus = analytical.min_gpus_analytical(
            lam, mu_gpu, t_slo_s, cv2, n_slots=n_slots
        )
        slo_optimal = True
    else:
        rho_c = max(0.05, min(0.95, rho_override))
        n_gpus = max(1, int(math.ceil(lam / (mu_gpu * rho_c))))
        slo_optimal = False

    rho = lam / (n_gpus * mu_gpu)
    n_active = rho * n_slots
    p99_ms = (
        analytical.p99_wait(n_gpus * n_slots, lam, mu_gpu / n_slots, cv2) + mean_pref
    ) * 1000.0

    try:
        pw = gpu.power_at_concurrency(max(1, int(n_active)))
    except ValueError:
        pw = float("nan")

    # per-pool tok/W (N cancels within the pool)
    tpw = mu_gpu * rho * mean_l_out / pw if (pw and pw > 0) else float("nan")
    cost = (n_gpus * gpu.cost_per_hr / 3600.0) / (lam * mean_l_out) * 1e6

    return TpwPoint(
        gpu_name=gpu.name,
        pool_label=pool_label or gpu.name,
        n_gpus=n_gpus,
        rho=rho,
        n_active=n_active,
        power_per_gpu_w=pw,
        tokens_per_watt=tpw,
        cost_per_mil=cost,
        p99_ttft_ms=p99_ms,
        power_model_qual=qual,
        slo_optimal=slo_optimal,
    )


def tpw_analysis(
    cdf: list,
    lam: float,
    gpus: list,
    t_slo_ms: float = 500.0,
    max_ctx: int = 8192,
    rho_sweep: list[float] | None = None,
) -> list[TpwPoint]:
    """Compute tokens-per-watt for a list of GPU profiles (single-pool each).

    All GPUs receive the **same** full CDF and arrival rate, so this function
    is only correct when comparing GPUs that serve the **same model** on the
    **same workload** — i.e. you are only changing the hardware, not the
    model.

    For multi-pool routing (homo vs hetero, or semantic-router small+large
    model), use :func:`fleet_tpw_analysis` instead, which computes the
    correct fleet-level aggregate across pools.

    .. warning::

        Pre-built profiles are calibrated for Llama-3-70B on 8-GPU TP (H100,
        A100) or 7B-class on single-GPU (A10G).  Comparing H100 (70B) vs
        A10G (7B) with this function mixes two different served models; the
        result reflects model-size difference, not GPU efficiency.  Set the
        same model by creating :class:`ManualProfile` instances manually
        before calling this function.

    Parameters
    ----------
    cdf       : workload CDF as list of (token_threshold, cumulative_frac).
    lam       : arrival rate (req/s).
    gpus      : GPU profiles to compare (list of ManualProfile).
    t_slo_ms  : P99 TTFT SLO (ms).
    max_ctx   : context window (tokens).
    rho_sweep : optional list of utilisation fractions at which to also
                compute tokens/watt (in addition to the SLO-optimal point).

    Returns
    -------
    List of TpwPoint.  SLO-optimal points have ``slo_optimal=True``.

    Notes
    -----
    **Do not over-interpret absolute values.** The tokens/watt numbers depend
    on power model quality (see ``power_model_qual`` field):
      - H100: HIGH (fitted to ML.ENERGY v3.0 measured data).
      - A100: FAIR (one measured anchor + FLOPS-scaling projection).
      - A10G: LOW  (projection only; no published batch-vs-power data).
    """
    points: list[TpwPoint] = []

    for gpu in gpus:
        pt = _tpw_one_pool(cdf, lam, gpu, t_slo_ms, max_ctx)
        points.append(pt)

        if rho_sweep:
            for rho in rho_sweep:
                pt_s = _tpw_one_pool(cdf, lam, gpu, t_slo_ms, max_ctx, rho_override=rho)
                points.append(pt_s)

    return points


def fleet_tpw_analysis(
    pools: list,
    lam_total: float,
    t_slo_ms: float = 500.0,
    topology: str = "",
) -> FleetTpwResult:
    """Compute fleet-level tokens-per-watt across a multi-pool routing topology.

    Unlike :func:`tpw_analysis`, N does **not** cancel across pools:
    ``fleet_tpw = Σ_i(λ_i × mean_L_out_i) / Σ_i(N_i × P_i(n_i))``

    Parameters
    ----------
    pools : list of dicts, one per pool.  Each dict must contain:
        ``gpu``     : ManualProfile for this pool.
        ``cdf``     : (sub-)CDF for this pool's traffic.  For a b_short split
                      use :func:`_split_cdf` to obtain per-pool sub-CDFs.
        ``lam``     : arrival rate into this pool (req/s).
        ``max_ctx`` : context window for this pool (tokens).
        ``label``   : human-readable label (e.g. "short-7B", "long-70B").
    lam_total : total arrival rate (req/s); used only for output annotation.
    t_slo_ms  : P99 TTFT SLO (ms) — applied to every pool independently.
    topology  : optional label for the routing topology.

    Returns
    -------
    :class:`FleetTpwResult` with per-pool :class:`TpwPoint` breakdowns and
    the fleet-level aggregate.

    Notes
    -----
    Each pool's CDF must be normalised (last cumulative fraction = 1.0).
    Use :func:`_split_cdf` to produce correctly normalised sub-CDFs.
    If any pool uses a LOW-quality power model, the fleet-level tok/W
    estimate inherits that uncertainty.
    """
    pool_points: list[TpwPoint] = []
    for spec in pools:
        pt = _tpw_one_pool(
            cdf=spec["cdf"],
            lam=spec["lam"],
            gpu=spec["gpu"],
            t_slo_ms=t_slo_ms,
            max_ctx=spec.get("max_ctx", 8192),
            pool_label=spec.get("label", spec["gpu"].name),
        )
        pool_points.append(pt)

    # Fleet-level aggregation: N does NOT cancel across pools
    total_output_tps = 0.0
    total_fleet_power_w = 0.0
    total_cost_per_s = 0.0
    total_gpus = 0
    worst_p99 = 0.0
    notes: list[str] = []

    for spec, pt in zip(pools, pool_points):
        mean_l_out = _cdf_mean_l_out(spec["cdf"])
        output_tps_pool = spec["lam"] * mean_l_out  # output tok/s from this pool
        pool_power_w = pt.n_gpus * pt.power_per_gpu_w  # W for whole pool

        total_output_tps += output_tps_pool
        total_fleet_power_w += pool_power_w
        total_cost_per_s += pt.n_gpus * spec["gpu"].cost_per_hr / 3600.0
        total_gpus += pt.n_gpus
        worst_p99 = max(worst_p99, pt.p99_ttft_ms)

        if pt.power_model_qual == "LOW":
            notes.append(
                f"  Pool '{pt.pool_label}' ({pt.gpu_name}): LOW power-model quality "
                f"— projection only, no published batch-vs-power measurements."
            )

    fleet_tpw = (
        total_output_tps / total_fleet_power_w
        if total_fleet_power_w > 0
        else float("nan")
    )
    fleet_cost = (total_cost_per_s / (lam_total * 1.0)) * 1e6
    # cost per 1M output tokens: total $/s / total output tok/s × 1e6
    fleet_cost_per_mil = (
        total_cost_per_s / total_output_tps * 1e6
        if total_output_tps > 0
        else float("nan")
    )

    if not topology:
        topology = " + ".join(p.pool_label for p in pool_points)

    return FleetTpwResult(
        topology=topology,
        pools=pool_points,
        fleet_tpw=fleet_tpw,
        fleet_cost_per_mil=fleet_cost_per_mil,
        fleet_power_kw=total_fleet_power_w / 1000.0,
        total_gpus=total_gpus,
        worst_p99_ms=worst_p99,
        power_model_notes=notes,
    )


def print_tpw_table(points: list[TpwPoint], title: str = "") -> None:
    """Print a formatted tokens-per-watt comparison table (single-pool entries)."""
    if title:
        print(f"\n{title}")
    hdr = (
        f"  {'Pool / GPU':22s}  {'GPUs':>5}  {'ρ':>5}  {'n_act':>6}  "
        f"{'Tok/W':>6}  {'P/GPU (W)':>9}  {'$/1M tok':>9}  "
        f"{'P99 (ms)':>9}  {'PwrQ':>4}  {'SLO':>3}"
    )
    sep = "  " + "-" * (len(hdr) - 2)
    print(sep)
    print(hdr)
    print(sep)
    for p in points:
        if not p.slo_optimal:
            continue
        label = p.pool_label or p.gpu_name
        flag = "★" if p.slo_optimal else " "
        qual_tag = p.power_model_qual
        if p.power_model_qual == "LOW":
            qual_tag += "*"
        print(
            f"  {label:22s}  {p.n_gpus:>5}  {p.rho:>5.2f}  "
            f"{p.n_active:>6.1f}  {p.tokens_per_watt:>6.2f}  "
            f"{p.power_per_gpu_w:>9.1f}  ${p.cost_per_mil:>8.2f}  "
            f"{p.p99_ttft_ms:>9.1f}  {qual_tag:>4}  {flag:>3}"
        )
    print(sep)
    if any(p.power_model_qual == "LOW" for p in points if p.slo_optimal):
        print(
            "  * LOW quality power model: projection only, no published measurements."
        )
        print("    Tok/W estimate for this GPU has high uncertainty.")
    print()


def print_fleet_tpw(result: FleetTpwResult, title: str = "") -> None:
    """Print a formatted fleet-level tokens-per-watt comparison."""
    if title:
        print(f"\n{title}")
    print(f"\n  Topology : {result.topology}")
    print(f"  Total GPUs     : {result.total_gpus}")
    print(f"  Fleet power    : {result.fleet_power_kw:.1f} kW")
    print(f"  Worst P99 TTFT : {result.worst_p99_ms:.1f} ms")
    print(
        f"  Fleet tok/W    : {result.fleet_tpw:.3f}  tok/J  ← correct multi-pool aggregate"
    )
    print(f"  Fleet $/1M tok : ${result.fleet_cost_per_mil:.2f}")

    hdr = (
        f"\n  {'Pool':22s}  {'GPU':18s}  {'GPUs':>5}  {'ρ':>5}  "
        f"{'Tok/W(pool)':>12}  {'P99(ms)':>9}  {'PwrQ':>4}"
    )
    print(hdr)
    print("  " + "-" * (len(hdr) - 3))
    for p in result.pools:
        label = p.pool_label or p.gpu_name
        qual_tag = p.power_model_qual + ("*" if p.power_model_qual == "LOW" else "")
        print(
            f"  {label:22s}  {p.gpu_name:18s}  {p.n_gpus:>5}  {p.rho:>5.2f}  "
            f"{p.tokens_per_watt:>12.3f}  {p.p99_ttft_ms:>9.1f}  {qual_tag:>4}"
        )
    print()
    if result.power_model_notes:
        print("  Power model caveats:")
        for note in result.power_model_notes:
            print(note)
        print()
