"""Replication (waste) accounting + optimization opportunities.

Two things this module produces:

1. `replication_report(cfg)` — where memory is being *duplicated* across
   the deployment, per category (attention weights, FFN weights, KV cache
   under replicate mode, DP replicas). Answers "if I lifted this
   duplication, how much would I save?"

2. `optimization_hints(cfg)` — a list of actionable hints. Each hint has
   a `category` (space | comm | compute), a `severity` (info | warn |
   good), and a short human message. UI code renders these as bullets.
"""
from __future__ import annotations

from dataclasses import dataclass

from .model_config import FullConfig
from .memory_layout import per_pe_weight_bytes, per_pe_kv_cache_bytes


# ── Replication accounting ────────────────────────────────────────

@dataclass
class ReplicationEntry:
    tensor: str
    per_pe_bytes: int
    replicated_across: str    # e.g. "CP (x4 copies)"
    copies: int               # total # copies globally
    wasted_bytes: int         # (copies - 1) * per_pe_bytes  (per PE view)


def _attn_weight_bytes_per_pe(cfg: FullConfig) -> int:
    """W_Q + W_K + W_V + W_O per PE, one layer."""
    m = cfg.model
    tp = cfg.topo.tp
    hq_per_pe = cfg.h_q_per_pe
    if cfg.topo.kv_shard_mode == "replicate":
        hkv_per_pe = max(1.0, m.h_kv / tp)
    else:
        hkv_per_pe = m.h_kv / tp
    per_layer = (
        m.hidden * hq_per_pe * m.d_head +
        m.hidden * hkv_per_pe * m.d_head +
        m.hidden * hkv_per_pe * m.d_head +
        hq_per_pe * m.d_head * m.hidden
    ) * m.bytes_per_elem
    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
    return int(per_layer * layers_per_stage)


def _ffn_weight_bytes_per_pe(cfg: FullConfig) -> int:
    """W_gate + W_up + W_down per PE, all local layers."""
    m = cfg.model
    ep = max(1, cfg.topo.ep)
    ffn_div = cfg.ffn_shard_divisor * ep
    per_layer = 3 * m.hidden * (m.ffn_dim // max(1, ffn_div)) * m.bytes_per_elem
    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
    return int(per_layer * layers_per_stage)


def _kv_bytes_per_pe(cfg: FullConfig) -> int:
    return per_pe_kv_cache_bytes(cfg)


def replication_report(cfg: FullConfig) -> list[ReplicationEntry]:
    """Enumerate replicated (duplicated) tensors and their waste."""
    entries: list[ReplicationEntry] = []
    topo = cfg.topo

    # Attention weights are NOT sharded by CP - so each CP rank holds a copy.
    attn_pe = _attn_weight_bytes_per_pe(cfg)
    cp_copies = topo.cp
    if cp_copies > 1:
        entries.append(ReplicationEntry(
            tensor="Attn weights (W_Q/W_K/W_V/W_O)",
            per_pe_bytes=attn_pe,
            replicated_across=f"CP (x{cp_copies} identical group copies)",
            copies=cp_copies,
            wasted_bytes=attn_pe * (cp_copies - 1),
        ))

    # FFN weights: replicated across scope levels NOT included.
    # Scope=TP  -> replicated across CP AND DP.
    # Scope=TP+CP -> replicated across DP only.
    # Scope=TP+CP+DP -> no replication (perfect share).
    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
    scope = topo.ffn_shard_scope
    ffn_cp_copies = 1 if "CP" in scope else topo.cp
    ffn_dp_copies = 1 if "DP" in scope else topo.dp
    ffn_total_copies = ffn_cp_copies * ffn_dp_copies
    if ffn_total_copies > 1:
        reasons = []
        if ffn_cp_copies > 1:
            reasons.append(f"CP (x{ffn_cp_copies})")
        if ffn_dp_copies > 1:
            reasons.append(f"DP (x{ffn_dp_copies})")
        entries.append(ReplicationEntry(
            tensor=f"FFN weights (scope={scope})",
            per_pe_bytes=ffn_pe,
            replicated_across=" & ".join(reasons),
            copies=ffn_total_copies,
            wasted_bytes=ffn_pe * (ffn_total_copies - 1),
        ))

    # KV cache: if kv_shard_mode='replicate' and TP > H_kv, each KV head
    # is duplicated tp // h_kv times across TP ranks.
    if (cfg.kv_replication_needed
            and topo.kv_shard_mode == "replicate"):
        kv_pe = _kv_bytes_per_pe(cfg)
        rep = cfg.head_dim_split_factor  # tp // h_kv
        entries.append(ReplicationEntry(
            tensor="KV cache (replicate mode)",
            per_pe_bytes=kv_pe,
            replicated_across=f"TP (x{rep} copies of each KV head)",
            copies=rep,
            wasted_bytes=kv_pe * (rep - 1),
        ))

    # DP replicas duplicate the entire per-PE footprint (attn + ffn + KV).
    if topo.dp > 1:
        total_pe = attn_pe + ffn_pe + _kv_bytes_per_pe(cfg)
        entries.append(ReplicationEntry(
            tensor="Full per-PE footprint (attn + FFN + KV)",
            per_pe_bytes=total_pe,
            replicated_across=f"DP (x{topo.dp} model replicas)",
            copies=topo.dp,
            wasted_bytes=total_pe * (topo.dp - 1),
        ))

    return entries


# ── Optimization hints ────────────────────────────────────────────

@dataclass
class Hint:
    category: str   # "space" | "comm" | "compute" | "layout"
    severity: str   # "info" | "warn" | "good"
    message: str


def optimization_hints(cfg: FullConfig) -> list[Hint]:
    hints: list[Hint] = []
    m = cfg.model
    topo = cfg.topo

    # SPACE hints -------------------------------------------------
    scope = topo.ffn_shard_scope
    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
    if scope == "TP" and topo.cp > 1:
        current = ffn_pe
        new = current // topo.cp
        hints.append(Hint(
            "space", "warn",
            f"FFN scope=TP replicates FFN weights across all {topo.cp} CP "
            f"groups. Switching to **TP+CP** would drop per-PE FFN weight "
            f"from {current/1e9:.2f} GB to {new/1e9:.2f} GB "
            f"(saves {(current-new)/1e9:.2f} GB/PE) at the cost of one "
            f"extra AllReduce per layer."
        ))
    if scope in ("TP", "TP+CP") and topo.dp > 1:
        base = ffn_pe if scope == "TP+CP" else ffn_pe // topo.cp
        new = base // topo.dp
        hints.append(Hint(
            "space", "info",
            f"With DP={topo.dp}, scope=**TP+CP+DP** would shard FFN across "
            f"replicas too, dropping per-PE FFN from {base/1e9:.2f} GB to "
            f"{new/1e9:.2f} GB (comes with an inter-replica AllReduce)."
        ))
    if (cfg.kv_replication_needed
            and topo.kv_shard_mode == "replicate"):
        rep = cfg.head_dim_split_factor
        kv_pe = _kv_bytes_per_pe(cfg)
        hints.append(Hint(
            "space", "warn",
            f"KV mode=replicate holds each KV head {rep}x. Switching to "
            f"**split** saves ~{kv_pe*(rep-1)/rep/1e9:.2f} GB KV per PE "
            f"but adds a Score AllReduce over {rep} ranks per hop."
        ))
    if topo.cp == 1 and topo.tp <= m.h_kv:
        # Suggest CP if s_kv is large and KV dominates.
        kv_pe = _kv_bytes_per_pe(cfg)
        attn_pe = _attn_weight_bytes_per_pe(cfg)
        if kv_pe > attn_pe:
            hints.append(Hint(
                "space", "info",
                f"KV cache ({kv_pe/1e9:.2f} GB) dominates weights "
                f"({attn_pe/1e9:.2f} GB) at S_kv={topo.s_kv:,}. Adding CP "
                f"shards the sequence axis (S_local = S_kv/CP)."
            ))

    # COMM hints --------------------------------------------------
    if topo.cp_inter_sip_hops > 0:
        hints.append(Hint(
            "comm", "warn",
            f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
            f"at {cfg.machine.bw_intersip_gbs:.0f} GB/s vs "
            f"{cfg.machine.bw_inter_gbs:.0f} GB/s intra-SIP - "
            f"consider a smaller CP or a topology that keeps CP inside one SIP."
        ))
    if topo.tp_spans_cubes > 1:
        hints.append(Hint(
            "comm", "warn",
            f"TP={topo.tp} spans {topo.tp_spans_cubes} cubes, so the W_O "
            f"AllReduce runs cross-cube "
            f"({cfg.machine.bw_intra_gbs:.0f} -> "
            f"{cfg.machine.bw_inter_gbs:.0f} GB/s). A TP that fits in one "
            f"cube (TP<={topo.pes_per_cube_hw}) keeps it intra-cube."
        ))
    if topo.sips_used > 1 and topo.sip_topology == "ring":
        hints.append(Hint(
            "layout", "info",
            f"With {topo.sips_used} SIPs on a ring, worst-case inter-SIP "
            f"distance is {topo.sips_used - 1} hops. Switching to **torus2d** "
            f"cuts worst-case hops to ~sqrt({topo.sips_used})."
        ))
    if topo.tp <= m.h_kv:
        hints.append(Hint(
            "comm", "good",
            f"TP ({topo.tp}) <= H_kv ({m.h_kv}): each KV head lives on "
            f"exactly one TP rank - no head-split AllReduce needed."
        ))

    # COMPUTE hints -----------------------------------------------
    if topo.mode == "decode" and topo.T_q == 1:
        # Decode is memory-bound almost everywhere; note it.
        hints.append(Hint(
            "compute", "info",
            "Decode (T_q=1): most GEMMs are memory-bound. FLOPs matter far "
            "less than HBM BW here - budget attention around "
            "weight-read cost, not compute peak."
        ))

    return hints