"""Model + machine parameters for the analytical visualization tool.

All numbers are per-rank / per-PE unless noted. bytes are counted in
bytes (b), FLOPs in raw ops, bandwidth in bytes/second, latency in
seconds. Convert to μs / GB / GFLOP at display time only.
"""
from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class ModelConfig:
    """Transformer model dimensions."""
    name: str = "Qwen 3 8B"
    hidden: int = 4096
    ffn_dim: int = 12288
    h_q: int = 32               # Q heads
    h_kv: int = 8               # KV heads
    d_head: int = 128
    layers: int = 36
    bytes_per_elem: int = 2     # bf16

    @property
    def group_size(self) -> int:
        """GQA group size = H_q / H_kv."""
        return self.h_q // self.h_kv

    @property
    def head_dim_total_q(self) -> int:
        return self.h_q * self.d_head

    @property
    def head_dim_total_kv(self) -> int:
        return self.h_kv * self.d_head


@dataclass
class TopologyConfig:
    """Physical mapping — DP × PP × CP × TP (× EP for MoE)."""
    cp: int = 4                 # inter-cube ring size (sequence shard)
    tp: int = 8                 # PEs per cube (head/hidden shard)
    pp: int = 1                 # pipeline stages (layer shard)
    dp: int = 1                 # data-parallel replicas
    ep: int = 1                 # expert-parallel degree (MoE only)
    s_kv: int = 4096
    mode: str = "decode"        # "decode" or "prefill"
    kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)
    # FFN sharding scope: default TP+CP (across cube group) trades a small
    # extra AllReduce for lower per-PE weight bytes.
    ffn_shard_scope: str = "TP+CP"   # "TP" | "TP+CP" | "TP+CP+DP"
    # Inter-SIP interconnect: how physical SIPs are wired together.
    # "ring" (1D wrap), "mesh2d" (grid, no wrap), "torus2d" (grid + wrap).
    sip_topology: str = "ring"
    # Placement of parallelism dims on the hierarchy.
    # "pe"   -> dim lives on PEs within a single cube (intra-cube).
    # "cube" -> dim lives across cubes (inter-cube).
    tp_placement: str = "pe"     # default: TP fills PEs of a cube
    cp_placement: str = "cube"   # default: CP fills cubes of a SIP
    # CP ring variant: what gets passed around the CP ring each hop.
    # "kv"    -> K and V shards rotate (Q,O,m,l stay local). Bytes/hop scales
    #            with S_local * H_kv * d_h. Good for prefill (large T_q).
    # "qoml"  -> Q and running (O, m, l) rotate; K,V stay local. Bytes/hop
    #            scales with T_q * H_q * d_h. Cheap for decode (T_q=1).
    cp_ring_variant: str = "kv"
    pes_per_cube_hw: int = 8    # SIP hardware constant
    cubes_per_sip_hw: int = 16

    @property
    def pes_per_stage(self) -> int:
        """PEs per pipeline stage (all stages have the same size)."""
        return self.cp * self.tp

    @property
    def pes_per_replica(self) -> int:
        """PEs per one full model replica (CP*TP*PP)."""
        return self.cp * self.tp * self.pp

    @property
    def total_pes(self) -> int:
        return self.pes_per_replica * self.dp

    # ── Placement-aware layout ────────────────────────────────
    @property
    def intra_cube_dims(self) -> int:
        """Product of dims placed on PE level (must fit in one cube, else spills)."""
        n = 1
        if self.tp_placement == "pe":
            n *= self.tp
        if self.cp_placement == "pe":
            n *= self.cp
        return n

    @property
    def inter_cube_dims(self) -> int:
        """Product of dims placed at cube level = distinct cube groups per stage."""
        n = 1
        if self.tp_placement == "cube":
            n *= self.tp
        if self.cp_placement == "cube":
            n *= self.cp
        return n

    @property
    def pes_per_cube_used(self) -> int:
        """How many PEs are live in each cube."""
        return min(self.intra_cube_dims, self.pes_per_cube_hw)

    @property
    def cubes_per_stage(self) -> int:
        """Cubes needed per PP stage under current placement.
        If intra-cube demand exceeds PEs/cube, the intra dim spills to cubes."""
        spill = max(1, (self.intra_cube_dims + self.pes_per_cube_hw - 1)
                    // self.pes_per_cube_hw)
        return self.inter_cube_dims * spill

    @property
    def cubes_used(self) -> int:
        """Physical cubes used across all PP stages and DP replicas."""
        return self.cubes_per_stage * self.pp * self.dp

    @property
    def sips_used(self) -> int:
        """Number of SIPs used (each SIP = pes_per_cube_hw × cubes_per_sip_hw)."""
        cubes_per_sip = self.cubes_per_sip_hw
        return max(1, (self.cubes_used + cubes_per_sip - 1) // cubes_per_sip)

    @property
    def placement_valid(self) -> bool:
        """False when intra-cube demand exceeds one cube (dim spills to cubes)."""
        return self.intra_cube_dims <= self.pes_per_cube_hw

    @property
    def tp_spans_cubes(self) -> int:
        """How many cubes one TP group spans (>=1). Depends on placement:
        - tp_placement=pe: 1 if TP fits in a cube, else spill.
        - tp_placement=cube: TP itself is the number of cubes per TP group.
        """
        if self.tp_placement == "cube":
            return max(1, self.tp)
        return max(1, (self.tp + self.pes_per_cube_hw - 1) // self.pes_per_cube_hw)

    @property
    def cp_intra_sip_hops(self) -> int:
        """CP ring hops that stay within one SIP."""
        if self.cp <= 1:
            return 0
        sips_in_ring = self.sips_used  # one PP stage's CP ring may span SIPs
        return max(0, (self.cp - 1) - (sips_in_ring - 1))

    @property
    def cp_inter_sip_hops(self) -> int:
        """CP ring hops that cross SIP boundaries."""
        if self.cp <= 1:
            return 0
        # Approximation: one inter-SIP hop per SIP boundary in the ring.
        # For CP=32 across 2 SIPs → 1 inter-SIP + 30 intra-SIP hops.
        return max(0, self.sips_used - 1)

    @property
    def T_q(self) -> int:
        return 1 if self.mode == "decode" else self.s_kv // self.cp

    @property
    def s_local(self) -> int:
        return self.s_kv // self.cp

    @property
    def layers_per_stage(self) -> str:
        """Layers per PP stage. Requires model.layers to render — leave as fn."""
        return "N_layers / PP"

    def tp_link_tier(self) -> str:
        """Return which BW tier a TP AllReduce uses: 'intra' | 'inter' | 'intersip'.
        - tp_placement=pe and TP fits in one cube -> 'intra'
        - tp_placement=cube and TP fits in one SIP -> 'inter'
        - otherwise -> 'intersip'
        """
        if self.tp_placement == "pe":
            if self.tp <= self.pes_per_cube_hw:
                return "intra"
            # spilled to multiple cubes: cross-cube
            if self.tp_spans_cubes <= self.cubes_per_sip_hw:
                return "inter"
            return "intersip"
        else:  # tp_placement == "cube"
            if self.tp <= self.cubes_per_sip_hw:
                return "inter"
            return "intersip"

    def cp_link_tier(self) -> str:
        """Return BW tier for the CP ring: 'intra' | 'inter' | 'intersip'."""
        if self.cp_placement == "pe":
            return "intra"
        # cp_placement == "cube"
        if self.cp <= self.cubes_per_sip_hw:
            return "inter"
        return "intersip"


@dataclass
class MachineParams:
    """SIP hardware parameters — all user-adjustable via sliders.

    Bandwidth tiers, from fastest to slowest:
      HBM (per-PE local memory)
        > intra-cube (PE↔PE on same die)
        > inter-cube (UCIe D2D between dies on same SIP)
        > inter-SIP (C2C / RDMA-class, across SIPs / servers)
    """
    # Per-PE compute peak (TFLOPs f16).
    peak_tflops_f16: float = 8.0
    # Per-PE HBM budget (GB) — the memory ceiling per PE.
    pe_hbm_gb: float = 6.0
    # HBM per-PE bandwidth (GB/s).
    bw_hbm_gbs: float = 256.0
    # Intra-cube PE↔PE link BW (GB/s).
    bw_intra_gbs: float = 512.0
    # Inter-cube (UCIe D2D) BW (GB/s).
    bw_inter_gbs: float = 128.0
    # Inter-SIP (C2C / RDMA) BW (GB/s).
    bw_intersip_gbs: float = 50.0
    # Per-hop latencies (ns).
    alpha_intra_ns: float = 20.0
    alpha_inter_ns: float = 100.0
    alpha_intersip_ns: float = 1000.0
    # Achievable utilization for compute-bound stages (0-1).
    compute_util: float = 0.8

    @property
    def peak_flops(self) -> float:
        return self.peak_tflops_f16 * 1e12

    @property
    def bw_hbm(self) -> float:
        return self.bw_hbm_gbs * 1e9

    @property
    def bw_intra(self) -> float:
        return self.bw_intra_gbs * 1e9

    @property
    def bw_inter(self) -> float:
        return self.bw_inter_gbs * 1e9

    @property
    def bw_intersip(self) -> float:
        return self.bw_intersip_gbs * 1e9

    @property
    def alpha_intra(self) -> float:
        return self.alpha_intra_ns * 1e-9

    @property
    def alpha_inter(self) -> float:
        return self.alpha_inter_ns * 1e-9

    @property
    def alpha_intersip(self) -> float:
        return self.alpha_intersip_ns * 1e-9

    @property
    def pe_budget_bytes(self) -> int:
        return int(self.pe_hbm_gb * 1e9)


@dataclass
class FullConfig:
    """Bundled config for one analysis run."""
    model: ModelConfig = field(default_factory=ModelConfig)
    topo: TopologyConfig = field(default_factory=TopologyConfig)
    machine: MachineParams = field(default_factory=MachineParams)

    @property
    def h_q_per_pe(self) -> int:
        return max(1, self.model.h_q // self.topo.tp)

    @property
    def h_kv_per_pe(self) -> float:
        """Fractional if TP > H_kv (head-dim split needed)."""
        return self.model.h_kv / self.topo.tp

    @property
    def d_per_pe(self) -> int:
        return self.model.hidden // self.topo.tp

    @property
    def ffn_per_pe(self) -> int:
        return self.model.ffn_dim // self.topo.tp

    @property
    def kv_replication_needed(self) -> bool:
        """True if TP > H_kv (each KV head shared across TP/H_kv ranks)."""
        return self.topo.tp > self.model.h_kv

    @property
    def head_dim_split_factor(self) -> int:
        """Number of TP ranks sharing one head (via d_head split)."""
        return max(1, self.topo.tp // self.model.h_kv)

    @property
    def ffn_shard_divisor(self) -> int:
        """How much the FFN dim is sharded across ranks."""
        scope = self.topo.ffn_shard_scope
        d = self.topo.tp
        if "CP" in scope:
            d *= self.topo.cp
        if "DP" in scope:
            d *= self.topo.dp
        return max(1, d)