"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers.

The PE-level view has always color-labeled CP ranks when they're packed
intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP
gets pushed to cube-level, so the intra-cube color branch is skipped —
which used to silently drop the CP rank locator entirely. These tests
guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are
always present when CP > 1, regardless of placement.
"""
from __future__ import annotations

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from tests.analytical_visualization.model_config import (
    FullConfig, MachineParams, TopologyConfig,
)
from tests.analytical_visualization.model_presets import PRESETS
from tests.analytical_visualization.pe_weight_layout import draw_pe_layout


def _cfg(cp: int, tp: int, cp_placement: str | None = None):
    model = PRESETS["Llama 3 8B"].model
    kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode")
    if cp_placement is not None:
        kwargs["cp_placement"] = cp_placement
    return FullConfig(
        model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(),
    )


def _all_text(fig) -> str:
    """Concatenate all text artists (labels + title) into one blob."""
    ax = fig.gca()
    parts = [t.get_text() for t in ax.texts]
    parts.append(ax.get_title())
    return "\n".join(parts)


def test_high_tp_cp_gt_1_shows_rank_locator():
    """TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator
    must still appear in cube tag + per-PE header + figure title."""
    cfg = _cfg(cp=4, tp=8)
    assert cfg.topo.cp_placement == "cube"  # sanity
    fig = draw_pe_layout(cfg)
    try:
        text = _all_text(fig)
        assert "CP rank 0 of 4" in text, text
        assert "CP=0" in text, text
        assert "CP=4" in text, text  # figure title dim summary
    finally:
        plt.close(fig)


def test_intra_cube_cp_placement_still_shows_per_rank_labels():
    """cp_placement='pe' (small TP, CP>1) — the old branch already
    colors + labels each CP rank. Regression check: the new tag does
    NOT replace those, and each CP=0..CP-1 label is still emitted."""
    cfg = _cfg(cp=4, tp=2)  # 4*2=8 fits intra-cube
    assert cfg.topo.cp_placement == "cube"  # default, but auto/manual can differ
    # Force pe placement to exercise the intra-cube branch.
    cfg.topo.cp_placement = "pe"
    fig = draw_pe_layout(cfg)
    try:
        text = _all_text(fig)
        for r in range(4):
            assert f"CP={r}" in text, (r, text)
        # The 'CP rank 0 of N' tag is only for the cube-placement fallback.
        assert "CP rank 0 of" not in text, text
    finally:
        plt.close(fig)


def test_cp_1_no_locator_added():
    """CP=1: nothing to locate; no CP tag or label should appear."""
    cfg = _cfg(cp=1, tp=8)
    fig = draw_pe_layout(cfg)
    try:
        text = _all_text(fig)
        assert "CP rank" not in text, text
        # The header 'CP=0' should not be added when CP=1.
        # (Figure title mentions 'CP=1' — that's an OK dim summary,
        # not a locator per-PE.)
        assert "  |  CP=0" not in text, text
    finally:
        plt.close(fig)


def test_tp_16_still_shows_cp_rank_across_spilled_cubes():
    """TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes
    for CP rank 0 must carry the '[CP rank 0 of 2]' tag."""
    cfg = _cfg(cp=2, tp=16)
    assert cfg.topo.cp_placement == "cube"
    fig = draw_pe_layout(cfg)
    try:
        text = _all_text(fig)
        assert text.count("CP rank 0 of 2") >= 2, text
        assert "CP=0" in text, text
    finally:
        plt.close(fig)