"""Data-parallel placement policy (ADR-0026: intra-device only).

``DPPolicy`` describes how a tensor is sharded *within a single SIP* across
that SIP's cubes and PEs. Crossing the SIP boundary is not a DPPolicy
concern: ADR-0024's ``torch.ahbm.set_device(rank)`` picks the SIP, and
Megatron-style TP (ADR-0027) expresses multi-SIP tensors when needed.

``ShardSpec`` is expressed in structural ``(sip, cube, pe)`` coordinates.
The former flat ``pe_index`` field/property is fully removed — callers
needing a flat integer key compute it explicitly at the call site.
"""
from __future__ import annotations

from dataclasses import dataclass
from math import ceil
from typing import Literal


@dataclass(frozen=True)
class DPPolicy:
    """Intra-device (cube × PE) data-parallel policy.

    SIP-level placement is controlled by ``torch.ahbm.set_device(rank)``
    (ADR-0024). For tensors that must cross SIP boundaries, use
    Megatron-style parallel layers (ADR-0027). DPPolicy itself never
    crosses a SIP boundary.

    Policies (per axis):
      - "replicate": full copy at each unit
      - "column_wise": split K (column) axis across units
      - "row_wise": split M (row) axis across units

    Optional overrides (``None`` = use topology dimensions):
      - num_pes: override PEs per cube
      - num_cubes: override cubes per SIP

    Cube range:
      - cube_start: index of the first cube in the SIP to place shards
        on. Defaults to 0. Enables disjoint sub-meshes within one SIP
        (e.g. cubes 8..15 alongside cubes 0..7) — required by the
        GQA Llama-70B 8-KV-group headline target (two 2×4 KV-groups
        per SIP).
    """

    cube: Literal["replicate", "column_wise", "row_wise"] = "replicate"
    pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
    num_pes: int | None = None
    num_cubes: int | None = None
    cube_start: int = 0


@dataclass(frozen=True)
class ShardSpec:
    """Structural shard placement — ``(sip, cube, pe)`` coord (ADR-0026).

    Global-flat ``pe_index`` was removed: callers must use structural
    coords directly. If a flat integer key is needed in a local context
    (e.g. internal dict lookup), compute it explicitly at the call site
    and do not expose it in any public API.
    """

    sip: int
    cube: int
    pe: int
    offset_bytes: int
    nbytes: int


@dataclass(frozen=True)
class _LocalPeShard:
    """Internal — PE resolver's return type (ADR-0026 D3).

    Holds a cube-local PE identifier (``local_pe``) plus the shard's
    byte payload. Lifted into ``ShardSpec`` with full ``(sip, cube, pe)``
    coordinates inside ``resolve_dp_policy``.
    """

    local_pe: int
    offset_bytes: int
    nbytes: int


def _split_shape(
    policy: str, shape: tuple[int, int], count: int, itemsize: int,
) -> list[tuple[tuple[int, int], int]]:
    """Split shape by policy into (sub_shape, byte_offset) pairs."""
    M, K = shape
    if policy == "replicate":
        return [((M, K), 0)] * count
    elif policy == "column_wise":
        chunk_k = K // count
        return [((M, chunk_k), i * M * chunk_k * itemsize) for i in range(count)]
    elif policy == "row_wise":
        chunk_m = M // count
        return [((chunk_m, K), i * chunk_m * K * itemsize) for i in range(count)]
    else:
        raise ValueError(f"Unknown policy: {policy}")


def resolve_dp_policy(
    policy: DPPolicy,
    *,
    shape: tuple[int, int],
    itemsize: int,
    num_pe: int,
    num_cubes: int = 1,
    target_sip: int,
) -> list[ShardSpec]:
    """Resolve a DPPolicy into a list[ShardSpec] on a single SIP.

    Two-level resolution (cube × PE) within ``target_sip``. Each returned
    ``ShardSpec`` carries ``sip=target_sip`` and cube/pe local to the SIP.
    No SIP-level split — DPPolicy is intra-device only (ADR-0026).
    """
    _PE_RESOLVERS = {
        "replicate": replicate,
        "column_wise": column_wise,
        "row_wise": row_wise,
    }
    resolver = _PE_RESOLVERS.get(policy.pe)
    if resolver is None:
        raise ValueError(f"Unknown pe-level policy: {policy.pe}")

    all_shards: list[ShardSpec] = []

    # Level 1: cube within SIP
    cube_splits = _split_shape(policy.cube, shape, num_cubes, itemsize)

    for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
        # Level 2: PE within cube — resolver returns _LocalPeShard
        local_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)

        for ls in local_shards:
            all_shards.append(ShardSpec(
                sip=target_sip,
                cube=policy.cube_start + cube_id,
                pe=ls.local_pe,
                offset_bytes=cube_offset + ls.offset_bytes,
                nbytes=ls.nbytes,
            ))

    return all_shards


def column_wise(
    *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[_LocalPeShard]:
    """Split K axis into num_pe equal parts. Each PE gets (M, K/P)."""
    M, K = shape
    chunk_k = K // num_pe
    chunk_bytes = M * chunk_k * itemsize
    return [
        _LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
        for i in range(num_pe)
    ]


def row_wise(
    *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[_LocalPeShard]:
    """Split M axis into num_pe equal parts. Each PE gets (M/P, K)."""
    M, K = shape
    chunk_m = M // num_pe
    chunk_bytes = chunk_m * K * itemsize
    return [
        _LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
        for i in range(num_pe)
    ]


def replicate(
    *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[_LocalPeShard]:
    """Full copy per PE. Each PE gets (M, K)."""
    M, K = shape
    full_bytes = M * K * itemsize
    return [
        _LocalPeShard(local_pe=i, offset_bytes=0, nbytes=full_bytes)
        for i in range(num_pe)
    ]


def tiled_column_major(
    *, shape: tuple[int, int], itemsize: int, num_pe: int,
    tile_m: int, tile_k: int,
) -> list[_LocalPeShard]:
    """2D tiling, column-major order (K axis first), round-robin across PEs."""
    M, K = shape
    tiles_m = ceil(M / tile_m)
    tiles_k = ceil(K / tile_k)
    tile_bytes = tile_m * tile_k * itemsize
    row_bytes = K * itemsize
    shards: list[_LocalPeShard] = []
    idx = 0
    for mi in range(tiles_m):
        for ki in range(tiles_k):
            offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
            shards.append(_LocalPeShard(
                local_pe=idx % num_pe,
                offset_bytes=offset,
                nbytes=tile_bytes,
            ))
            idx += 1
    return shards


def tiled_row_major(
    *, shape: tuple[int, int], itemsize: int, num_pe: int,
    tile_m: int, tile_k: int,
) -> list[_LocalPeShard]:
    """2D tiling, row-major order (M axis first), round-robin across PEs."""
    M, K = shape
    tiles_m = ceil(M / tile_m)
    tiles_k = ceil(K / tile_k)
    tile_bytes = tile_m * tile_k * itemsize
    row_bytes = K * itemsize
    shards: list[_LocalPeShard] = []
    idx = 0
    for ki in range(tiles_k):
        for mi in range(tiles_m):
            offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
            shards.append(_LocalPeShard(
                local_pe=idx % num_pe,
                offset_bytes=offset,
                nbytes=tile_bytes,
            ))
            idx += 1
    return shards