"""SFR configuration for the full IPCQ hardware wiring.

Installs PE_IPCQ neighbor tables modeling the physical hardware.
Wiring is independent of DPPolicy / kernel choice — the kernel decides
at runtime which links to use.

Direction label namespaces (disjoint):

  - Intra-cube PE-to-PE:   ``intra_N / intra_S / intra_E / intra_W``
    Logical 2×4 PE grid within a cube (no wrap):

         Row 0:  pe0  pe1  pe2  pe3
         Row 1:  pe4  pe5  pe6  pe7

  - Intercube same-lane:   ``N / S / E / W``
    ``pe_i of cube_A ↔ pe_i of cube_B`` across the 4×4 cube mesh
    (no wrap). Every PE i ∈ [0..7] wired independently.

  - Inter-SIP same-(cube, pe): ``global_N / global_S / global_E / global_W``
    ``pe_i of cube_c on sip_A ↔ pe_i of cube_c on sip_B`` per
    ``topology.yaml → system.sips.topology``.
"""
from __future__ import annotations

import types
from typing import Any

from kernbench.ccl.install import install_ipcq
from kernbench.ccl.topologies import _BUILTIN as _TOPO_BUILTINS


# ── Intra-cube 2×4 PE grid ───────────────────────────────────────────

_PE_GRID_COLS = 4
_PE_GRID_ROWS = 2
_PES_PER_CUBE = _PE_GRID_COLS * _PE_GRID_ROWS  # 8


def _intra_cube_neighbors(pe: int) -> dict[str, int]:
    """Logical 2×4 PE grid neighbors within a cube (no wrap).

    Returns directions in the ``intra_*`` namespace.
    """
    row, col = divmod(pe, _PE_GRID_COLS)
    nbrs: dict[str, int] = {}
    if col < _PE_GRID_COLS - 1:
        nbrs["intra_E"] = row * _PE_GRID_COLS + (col + 1)
    if col > 0:
        nbrs["intra_W"] = row * _PE_GRID_COLS + (col - 1)
    if row < _PE_GRID_ROWS - 1:
        nbrs["intra_S"] = (row + 1) * _PE_GRID_COLS + col
    if row > 0:
        nbrs["intra_N"] = (row - 1) * _PE_GRID_COLS + col
    return nbrs


# ── Public entry point ───────────────────────────────────────────────


def configure_sfr_intercube_multisip(
    engine: Any,
    spec: dict,
    cfg: dict,
) -> dict[str, Any]:
    """Wire the full IPCQ hardware model.

    Every PE on every cube on every SIP gets neighbor table entries for:

      - intra-cube (2×4 grid) in the ``intra_*`` namespace
      - intercube same-lane (4×4 cube mesh, no wrap) in ``N/S/E/W``
      - inter-SIP same-(cube, pe) in ``global_*``

    Args:
        engine: GraphEngine with ``_components``.
        spec: topology spec dict (from topology.yaml).
        cfg: merged algorithm config (from ``resolve_algorithm_config``).

    Returns:
        The install plan dict from ``install_ipcq``.
    """
    cm = spec["sip"]["cube_mesh"]
    mesh_w = int(cm["w"])
    mesh_h = int(cm["h"])
    n_cubes = mesh_w * mesh_h
    sips_cfg = spec.get("system", {}).get("sips", {})
    n_sips = int(sips_cfg.get("count", 1))
    sip_topology = str(sips_cfg.get("topology", "ring_1d"))
    sip_w = sips_cfg.get("w")
    sip_h = sips_cfg.get("h")
    sip_w = int(sip_w) if sip_w is not None else None
    sip_h = int(sip_h) if sip_h is not None else None

    if sip_topology not in _TOPO_BUILTINS:
        raise ValueError(
            f"Unknown sip topology '{sip_topology}'. "
            f"Available: {list(_TOPO_BUILTINS)}"
        )
    _sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology]

    def sip_topo_fn(rank: int, ws: int) -> dict:
        if sip_w is not None and sip_h is not None:
            try:
                return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h)
            except TypeError:
                pass
        return _sip_topo_fn_raw(rank, ws)

    pes_per_cube = _PES_PER_CUBE
    world_size = n_sips * n_cubes * pes_per_cube
    pe_idx_to_pe: list[tuple[int, int, int]] = [
        (sip, cube, pe)
        for sip in range(n_sips)
        for cube in range(n_cubes)
        for pe in range(pes_per_cube)
    ]

    def _pe_idx(sip: int, cube: int, pe: int) -> int:
        return (sip * n_cubes + cube) * pes_per_cube + pe

    def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]:
        tmp = pe_idx
        pe = tmp % pes_per_cube
        tmp //= pes_per_cube
        cube = tmp % n_cubes
        sip = tmp // n_cubes
        row = cube // mesh_w
        col = cube % mesh_w

        nbrs: dict[str, int] = {}

        # ── Intra-cube (intra_N/S/E/W) ──
        for d, peer_pe in _intra_cube_neighbors(pe).items():
            nbrs[d] = _pe_idx(sip, cube, peer_pe)

        # ── Intercube same-lane (N/S/E/W, 4×4 no wrap) ──
        if col < mesh_w - 1:
            nbrs["E"] = _pe_idx(sip, row * mesh_w + (col + 1), pe)
        if col > 0:
            nbrs["W"] = _pe_idx(sip, row * mesh_w + (col - 1), pe)
        if row < mesh_h - 1:
            nbrs["S"] = _pe_idx(sip, (row + 1) * mesh_w + col, pe)
        if row > 0:
            nbrs["N"] = _pe_idx(sip, (row - 1) * mesh_w + col, pe)

        # ── Inter-SIP same-(cube, pe) (global_*) ──
        if n_sips > 1:
            sip_nbrs = sip_topo_fn(sip, n_sips)
            for d, peer_sip in sip_nbrs.items():
                nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe)

        return nbrs

    mock_module = types.SimpleNamespace(neighbors=_neighbors)

    cfg_copy = dict(cfg)
    cfg_copy["world_size"] = world_size
    cfg_copy["topology"] = "none"

    return install_ipcq(
        engine, spec, cfg_copy,
        algo_module=mock_module,
        rank_to_pe=pe_idx_to_pe,
    )
