"""Auto-explore parallelism configuration space and rank by Pareto frontier.

Extends the existing ``autosuggest`` (memory-only, single winner) to the
full 9-knob search (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
tp_placement, cp_placement, cp_ring_variant) with four objectives:

    - min total_latency_ns   (single-request decode/prefill step)
    - max throughput_tok_s   (naive tokens/sec = 1 / latency for single-req)
    - max efficiency_score   (geo-mean of compute + BW utilization)
    - min pes_used

Reuses ``stage_latencies`` and ``memory_layout`` as the physics; adds
enumeration + 4D Pareto sort on top.

Analytical model runs in ~microseconds per config, so the full sweep
(~5-10k feasible configs after pruning) completes in a few seconds.
"""
from __future__ import annotations

import math
from collections.abc import Iterator
from dataclasses import dataclass, field, replace

from .memory_layout import compute_memory
from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
from .stage_latencies import all_ffn_stages, all_stages

# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity)
# Multiples of 2 (finer grid than powers of 2 alone), capped at values that
# are physically meaningful for the modelled topology.
_PARALLELISM_SWEEP_VALUES = {
    "cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256),
    "tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64),
    "pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32),
    "dp": (1, 2, 4, 6, 8, 12, 16),
    "ep": (1, 2, 4, 6, 8, 12, 16, 32),
}


# ── Search space ─────────────────────────────────────────────────────

_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
_PP_OPTIONS = (1, 2, 4, 8, 16)
_DP_OPTIONS = (1, 2, 4)
_KV_SHARD_MODES = ("split", "replicate")
_FFN_SHARD_SCOPES = ("TP", "TP+CP", "TP+CP+DP")
_TP_PLACEMENTS = ("pe", "cube")
_CP_PLACEMENTS = ("cube", "pe")
_CP_RING_VARIANTS = ("kv", "qoml")


# ── Result types ─────────────────────────────────────────────────────


@dataclass
class ConfigScore:
    """A single (config, computed-metrics) tuple. All floats in SI units
    (seconds, tokens/sec, dimensionless) unless suffixed otherwise."""

    # ── Config (the 9 knobs) ─────
    cp: int
    tp: int
    pp: int
    dp: int
    kv_shard_mode: str
    ffn_shard_scope: str
    tp_placement: str
    cp_placement: str
    cp_ring_variant: str

    # ── Objectives ─────
    total_latency_ns: float      # ↓ minimize
    throughput_tok_s: float       # ↑ maximize
    efficiency_score: float       # ↑ maximize (0..1)
    pes_used: int                 # ↓ minimize

    # ── Info-only (not ranked on) ─────
    hbm_utilization: float        # bytes_used / budget (0..1+; may exceed 1 if over-budget)
    weights_gb: float
    kv_gb: float
    transient_gb: float
    sips_used: int
    fits_memory: bool
    placement_valid: bool
    reason: str = ""

    @property
    def latency_us(self) -> float:
        return self.total_latency_ns / 1e3

    @property
    def latency_ms(self) -> float:
        return self.total_latency_ns / 1e6

    def as_topology(self, s_kv: int, mode: str) -> TopologyConfig:
        """Reconstruct the TopologyConfig this score was computed for."""
        return TopologyConfig(
            cp=self.cp, tp=self.tp, pp=self.pp, dp=self.dp,
            s_kv=s_kv, mode=mode,
            kv_shard_mode=self.kv_shard_mode,
            ffn_shard_scope=self.ffn_shard_scope,
            tp_placement=self.tp_placement,
            cp_placement=self.cp_placement,
            cp_ring_variant=self.cp_ring_variant,
        )


@dataclass
class AutoExploreResult:
    model_name: str
    s_kv: int
    mode: str
    total_enumerated: int         # after basic domain pruning
    total_feasible: int           # after memory + placement checks
    all_scores: list[ConfigScore] = field(default_factory=list)  # every feasible config
    pareto_scores: list[ConfigScore] = field(default_factory=list)  # non-dominated set


# ── Enumeration + pruning ────────────────────────────────────────────


def enumerate_configs(
    model: ModelConfig,
    s_kv: int,
    mode: str,
) -> Iterator[TopologyConfig]:
    """Yield every domain-valid TopologyConfig.

    Domain rules (fast pruning; no memory check yet):
      - PP ≤ model.layers                            (can't have more stages than layers)
      - TP ≤ 4 × model.h_q                            (unrealistic head-dim splits above this)
      - Skip ffn_shard_scope containing 'DP' when DP=1 (redundant with plain 'TP+CP')
      - Skip cp_ring_variant='qoml' when CP=1        (no ring, variant is a no-op)
    """
    for cp in _CP_OPTIONS:
        for tp in _TP_OPTIONS:
            if tp > 4 * model.h_q:
                continue
            for pp in _PP_OPTIONS:
                if pp > model.layers:
                    continue
                for dp in _DP_OPTIONS:
                    for kv_mode in _KV_SHARD_MODES:
                        for ffn_scope in _FFN_SHARD_SCOPES:
                            if "DP" in ffn_scope and dp == 1:
                                continue
                            for tp_place in _TP_PLACEMENTS:
                                for cp_place in _CP_PLACEMENTS:
                                    for cp_ring in _CP_RING_VARIANTS:
                                        if cp == 1 and cp_ring == "qoml":
                                            continue
                                        yield TopologyConfig(
                                            cp=cp, tp=tp, pp=pp, dp=dp,
                                            s_kv=s_kv, mode=mode,
                                            kv_shard_mode=kv_mode,
                                            ffn_shard_scope=ffn_scope,
                                            tp_placement=tp_place,
                                            cp_placement=cp_place,
                                            cp_ring_variant=cp_ring,
                                        )


# ── Scoring ──────────────────────────────────────────────────────────


def _sum_visible_latency(
    cfg: FullConfig,
    include_attention: bool = True,
    include_ffn: bool = True,
) -> float:
    """Total single-request latency (seconds) across all model layers.

    A single request traverses every layer sequentially, whether the layers
    sit on one PP stage or are spread across many:
      - PP=1  : one rank holds all L layers, latency = L × per_layer
      - PP=K  : K ranks each hold L/K layers, but request crosses all K
                stages sequentially → still L × per_layer

    PP therefore does NOT reduce single-request latency; it only improves
    throughput under batching. This function is the single-request cost, so
    we multiply by full model.layers regardless of PP.

    Scope selection (both default True — full per-token cost):
      - ``include_attention=True,  include_ffn=True``  → full transformer
      - ``include_attention=True,  include_ffn=False`` → attention only
      - ``include_attention=False, include_ffn=True``  → FFN / MoE only
      - ``include_attention=False, include_ffn=False`` → zero (rejected caller-side)
    """
    attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0
    ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0
    per_layer = attn + ffn
    return per_layer * cfg.model.layers


def _efficiency(cfg: FullConfig, latency_s: float,
                include_attention: bool = True,
                include_ffn: bool = True) -> float:
    """Geo-mean of compute-util and BW-util. Range ~ (0, 1].

    - compute_util = achieved_flops / (peak_flops × pes × latency)
    - bw_util      = achieved_bytes / (peak_bw × pes × latency)

    Scope flags mirror :func:`_sum_visible_latency`.
    """
    if latency_s <= 0:
        return 0.0
    attn = all_stages(cfg) if include_attention else []
    ffn = all_ffn_stages(cfg) if include_ffn else []
    layers = math.ceil(cfg.model.layers / cfg.topo.pp)
    total_flops = layers * sum(s.flops for s in attn + ffn)
    total_bytes = layers * sum(s.mem_bytes for s in attn + ffn)

    pes = cfg.topo.total_pes
    peak_flops = cfg.machine.peak_flops * pes
    peak_bw = cfg.machine.bw_hbm * pes

    compute_util = total_flops / (peak_flops * latency_s) if peak_flops > 0 else 0.0
    bw_util = total_bytes / (peak_bw * latency_s) if peak_bw > 0 else 0.0
    compute_util = min(1.0, max(0.0, compute_util))
    bw_util = min(1.0, max(0.0, bw_util))
    # Geo-mean; if either is 0 the score is 0 (avoids overrewarding lopsided configs).
    return math.sqrt(compute_util * bw_util)


def score_config(cfg: FullConfig,
                 include_attention: bool = True,
                 include_ffn: bool = True) -> ConfigScore:
    """Compute all 4 objectives + info fields for one config.

    Feasibility (memory + placement) is stored but does NOT gate scoring —
    infeasible configs get returned with fits_memory=False so callers can
    filter or display them.

    Scope flags (default: full transformer) restrict *latency + efficiency*
    to attention only, FFN only, or both. Memory feasibility is unchanged —
    still checks weights+KV+transient fit in per-PE HBM since the model
    still exists physically regardless of what the caller is scoring.
    """
    mem = compute_memory(cfg)
    placement_ok = cfg.topo.placement_valid

    latency_s = _sum_visible_latency(
        cfg, include_attention=include_attention, include_ffn=include_ffn,
    )
    throughput = 1.0 / latency_s if latency_s > 0 else 0.0
    efficiency = (
        _efficiency(cfg, latency_s,
                    include_attention=include_attention, include_ffn=include_ffn)
        if latency_s > 0 else 0.0
    )
    fits = not mem.over_budget

    reason = ""
    if not fits:
        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
                  f"exceeds per-PE budget ({mem.budget_bytes/1e9:.2f} GB)")
    elif not placement_ok:
        reason = (f"intra-cube demand ({cfg.topo.intra_cube_dims}) "
                  f"exceeds PEs/cube ({cfg.topo.pes_per_cube_hw})")

    return ConfigScore(
        cp=cfg.topo.cp, tp=cfg.topo.tp, pp=cfg.topo.pp, dp=cfg.topo.dp,
        kv_shard_mode=cfg.topo.kv_shard_mode,
        ffn_shard_scope=cfg.topo.ffn_shard_scope,
        tp_placement=cfg.topo.tp_placement,
        cp_placement=cfg.topo.cp_placement,
        cp_ring_variant=cfg.topo.cp_ring_variant,
        total_latency_ns=latency_s * 1e9,
        throughput_tok_s=throughput,
        efficiency_score=efficiency,
        pes_used=cfg.topo.total_pes,
        hbm_utilization=mem.used_bytes / mem.budget_bytes if mem.budget_bytes > 0 else 0.0,
        weights_gb=mem.weights_bytes / 1e9,
        kv_gb=mem.kv_cache_bytes / 1e9,
        transient_gb=mem.transient_bytes / 1e9,
        sips_used=cfg.topo.sips_used,
        fits_memory=fits,
        placement_valid=placement_ok,
        reason=reason,
    )


# ── Pareto sort ──────────────────────────────────────────────────────


def _dominates(a: ConfigScore, b: ConfigScore) -> bool:
    """True iff a is no worse than b on every axis AND strictly better on ≥ 1.

    Axes (with direction) — 3D Pareto:
      total_latency_ns    ↓  (a ≤ b)
      pes_used            ↓  (a ≤ b)  — proxy for cost / deployment size
      efficiency_score    ↑  (a ≥ b)  — geo-mean of compute+BW utilization

    Note: ``throughput_tok_s`` is deliberately NOT a Pareto axis. For a
    single-request analysis, throughput = 1 / latency, so it's collinear
    with latency and would collapse the frontier. It stays on ConfigScore
    for display but doesn't participate in domination.
    """
    no_worse = (
        a.total_latency_ns <= b.total_latency_ns
        and a.pes_used <= b.pes_used
        and a.efficiency_score >= b.efficiency_score
    )
    if not no_worse:
        return False
    strictly_better = (
        a.total_latency_ns < b.total_latency_ns
        or a.pes_used < b.pes_used
        or a.efficiency_score > b.efficiency_score
    )
    return strictly_better


def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
    """Extract non-dominated set on (latency↓, pe↓, throughput↑, efficiency↑).

    Only feasible configs (fits_memory AND placement_valid) participate; the
    non-feasible are excluded from the frontier (they can still be returned in
    ``all_scores`` for informational display).

    O(N²) — fine for N up to ~50k with early exits.
    """
    feasible = [s for s in scores if s.fits_memory and s.placement_valid]
    frontier: list[ConfigScore] = []
    for i, a in enumerate(feasible):
        dominated = False
        for j, b in enumerate(feasible):
            if i == j:
                continue
            if _dominates(b, a):
                dominated = True
                break
        if not dominated:
            frontier.append(a)
    return frontier


# ── Parallelism sensitivity ─────────────────────────────────────────


@dataclass
class ParallelismSensitivityRow:
    """One knob's sweep result: latency + fit at every tested value.

    All *other* parallelism knobs (including HW) are held at their baseline
    values, so this isolates the effect of just this one knob.
    """
    knob: str                       # "cp" / "tp" / "pp" / "dp" / "ep"
    values: list[int]               # tested values (from _PARALLELISM_SWEEP_VALUES)
    latencies_ns: list[float]       # one per value; NaN when infeasible
    fits_flags: list[bool]          # per-value memory+placement feasibility
    baseline_value: int             # what the baseline config uses for this knob
    baseline_latency_ns: float


def compute_parallelism_sensitivity(
    baseline: ConfigScore,
    model: ModelConfig,
    machine: MachineParams,
    s_kv: int,
    mode: str,
    include_attention: bool = True,
    include_ffn: bool = True,
) -> list[ParallelismSensitivityRow]:
    """For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
    holding the OTHER knobs fixed at the baseline. Reports latency + memory
    fit per value.

    Answers: 'if I only change this one knob, how does latency and memory
    fit change?' — the complement to auto_hardware.compute_sensitivity
    which does the same for hardware knobs.
    """
    baseline_topo = baseline.as_topology(s_kv, mode)

    rows: list[ParallelismSensitivityRow] = []
    for knob, values in _PARALLELISM_SWEEP_VALUES.items():
        baseline_val = getattr(baseline_topo, knob, 1)
        # EP isn't stored on ConfigScore's as_topology output today (dp
        # attribute exists but ep does not always). Fall back to 1.
        latencies: list[float] = []
        fits: list[bool] = []
        for v in values:
            # Skip nonsensical values that would violate hard bounds.
            if knob == "pp" and v > model.layers:
                latencies.append(float("nan"))
                fits.append(False)
                continue
            if knob == "tp" and v > 4 * model.h_q:
                latencies.append(float("nan"))
                fits.append(False)
                continue
            # Build a variant TopologyConfig with just this one knob changed.
            trial = replace(baseline_topo, **{knob: v})
            cfg = FullConfig(model=model, topo=trial, machine=machine)
            score = score_config(
                cfg, include_attention=include_attention, include_ffn=include_ffn,
            )
            latencies.append(score.total_latency_ns)
            fits.append(score.fits_memory and score.placement_valid)
        rows.append(ParallelismSensitivityRow(
            knob=knob,
            values=list(values),
            latencies_ns=latencies,
            fits_flags=fits,
            baseline_value=int(baseline_val),
            baseline_latency_ns=baseline.total_latency_ns,
        ))
    return rows


# ── Top-level driver ─────────────────────────────────────────────────


def run_auto_explore(
    model: ModelConfig,
    machine: MachineParams,
    s_kv: int,
    mode: str = "decode",
    include_attention: bool = True,
    include_ffn: bool = True,
) -> AutoExploreResult:
    """Enumerate all configs, score each, extract Pareto frontier.

    Returns both the full ``all_scores`` list (for the table view) and the
    ``pareto_scores`` subset (for the scatter/highlight view). Both are sorted
    by total_latency_ns ascending.

    Scope: at least one of ``include_attention`` / ``include_ffn`` must be
    True. Setting both False would give zero latency for every config —
    the caller should not do that.
    """
    all_scores: list[ConfigScore] = []
    total_enumerated = 0
    for topo in enumerate_configs(model, s_kv, mode):
        total_enumerated += 1
        cfg = FullConfig(model=model, topo=topo, machine=machine)
        all_scores.append(score_config(
            cfg, include_attention=include_attention, include_ffn=include_ffn,
        ))

    all_scores.sort(key=lambda s: s.total_latency_ns)
    pareto = pareto_frontier(all_scores)
    pareto.sort(key=lambda s: s.total_latency_ns)

    feasible_count = sum(1 for s in all_scores if s.fits_memory and s.placement_valid)

    return AutoExploreResult(
        model_name=model.name,
        s_kv=s_kv,
        mode=mode,
        total_enumerated=total_enumerated,
        total_feasible=feasible_count,
        all_scores=all_scores,
        pareto_scores=pareto,
    )
