"""Phase 1 spec test for ``DPPolicy.cube_start``.

``cube_start`` is an optional override that shifts the cube indices in
``ShardSpec`` generated by ``resolve_dp_policy``. It enables addressing a
disjoint cube sub-mesh within one SIP (e.g. cubes 8..15 in addition to
0..7), which the GQA Llama-70B 8-KV-group headline target requires
(two 2×4 KV-groups per SIP × 4 SIPs = 64 cubes).

Default ``cube_start=0`` preserves every existing call-site bit-for-bit:
``DPPolicy(num_cubes=8)`` still resolves to ``ShardSpec.cube ∈ [0, 8)``.
The CCL milestone bench's full-SIP ``DPPolicy(num_cubes=16)`` likewise
continues to produce cubes 0..15.

Phase 1: production code is unchanged → these tests FAIL until the
Phase 2 diff lands. Phase 2 makes all of them pass.
"""
from __future__ import annotations

from kernbench.policy.placement.dp import DPPolicy, resolve_dp_policy


# ── Default ``cube_start=0`` preserves existing behavior ─────────────


def test_dppolicy_default_cube_start_is_zero():
    """``DPPolicy(num_cubes=8)`` sets ``cube_start=0`` by default."""
    dp = DPPolicy(cube="row_wise", num_cubes=8)
    assert dp.cube_start == 0


def test_resolve_default_yields_cubes_zero_through_seven():
    """Backward-compat: default ``cube_start=0`` → cubes 0..num_cubes-1."""
    dp = DPPolicy(cube="row_wise", pe="replicate", num_cubes=8, num_pes=1)
    shards = resolve_dp_policy(
        dp, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    cube_ids = sorted({s.cube for s in shards})
    assert cube_ids == list(range(8))


def test_resolve_full_sip_default_yields_cubes_zero_through_fifteen():
    """CCL pattern: ``num_cubes=16, cube_start=0`` → cubes 0..15 unchanged."""
    dp = DPPolicy(cube="row_wise", pe="replicate", num_cubes=16, num_pes=1)
    shards = resolve_dp_policy(
        dp, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=16, target_sip=0,
    )
    cube_ids = sorted({s.cube for s in shards})
    assert cube_ids == list(range(16))


# ── ``cube_start=N`` shifts cube indices ─────────────────────────────


def test_resolve_cube_start_eight_yields_cubes_eight_through_fifteen():
    """``cube_start=8, num_cubes=8`` → ``ShardSpec.cube ∈ [8, 16)``.

    This is the headline-enabling case: the second 2×4 KV-group per SIP
    in the 8-KV-group Llama-70B target lands on cubes 8..15.
    """
    dp = DPPolicy(
        cube="row_wise", pe="replicate",
        num_cubes=8, num_pes=1, cube_start=8,
    )
    shards = resolve_dp_policy(
        dp, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    cube_ids = sorted({s.cube for s in shards})
    assert cube_ids == list(range(8, 16))


def test_resolve_cube_start_preserves_shard_count():
    """Shifting ``cube_start`` does not change the total shard count."""
    dp_default = DPPolicy(
        cube="row_wise", pe="replicate", num_cubes=8, num_pes=1,
    )
    dp_shifted = DPPolicy(
        cube="row_wise", pe="replicate",
        num_cubes=8, num_pes=1, cube_start=8,
    )
    shards_def = resolve_dp_policy(
        dp_default, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    shards_shf = resolve_dp_policy(
        dp_shifted, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    assert len(shards_def) == len(shards_shf)


# ── ``cube_start`` works with ``cube="replicate"`` too ──────────────


def test_resolve_cube_start_with_replicate_policy():
    """``cube="replicate", cube_start=8`` → every cube in [8, 16) gets a copy."""
    dp = DPPolicy(
        cube="replicate", pe="replicate",
        num_cubes=8, num_pes=1, cube_start=8,
    )
    shards = resolve_dp_policy(
        dp, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    cube_ids = sorted({s.cube for s in shards})
    assert cube_ids == list(range(8, 16))


# ── No duplicates regardless of ``cube_start`` ───────────────────────


def test_resolve_cube_start_yields_distinct_cube_ids():
    """For ``pe="replicate", num_pes=1``, every shard has a distinct cube."""
    dp = DPPolicy(
        cube="row_wise", pe="replicate",
        num_cubes=8, num_pes=1, cube_start=8,
    )
    shards = resolve_dp_policy(
        dp, shape=(32, 64), itemsize=2,
        num_pe=1, num_cubes=8, target_sip=0,
    )
    cube_ids = [s.cube for s in shards]
    assert len(cube_ids) == len(set(cube_ids))
    assert all(8 <= c < 16 for c in cube_ids)