import pytest
from giant.cond_layout import AXIS_TYPES, CondLayout
from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM

# ── cond_cat column layout ───────────────────────────────────────────────────


def test_topn_cols_neither_onehot():
    layout = CondLayout.from_types("physical", "embedding")
    assert (layout.particle_topn_col, layout.material_topn_col) == (None, None)
    assert layout.cat_dim == 2


def test_topn_cols_particle_only():
    layout = CondLayout.from_types("onehot", "physical")
    assert (layout.particle_topn_col, layout.material_topn_col) == (2, None)
    assert layout.cat_dim == 3


def test_topn_cols_material_only():
    layout = CondLayout.from_types("physical", "onehot")
    assert (layout.particle_topn_col, layout.material_topn_col) == (None, 2)
    assert layout.cat_dim == 3


def test_topn_cols_both_onehot_particle_then_material():
    layout = CondLayout.from_types("onehot", "onehot")
    assert (layout.particle_topn_col, layout.material_topn_col) == (2, 3)
    assert layout.cat_dim == 4


def test_dense_vocab_cols_are_mode_independent():
    """Columns 0/1 are always the dense pdg/material index — giant.model.routers
    reads them without knowing the conditioning mode."""
    assert (CondLayout.PDG_COL, CondLayout.MAT_COL) == (0, 1)
    for particle in AXIS_TYPES:
        for material in AXIS_TYPES:
            layout = CondLayout.from_types(particle, material)
            assert layout.particle_topn_col not in (layout.PDG_COL, layout.MAT_COL)
            assert layout.material_topn_col not in (layout.PDG_COL, layout.MAT_COL)


# ── cond_cont slice layout ───────────────────────────────────────────────────


def test_cont_slices_tile_cond_cont_exactly():
    """base / particle_phys / material_phys must partition cond_cont with no
    gap and no overlap — a gap or overlap is exactly the silent
    mis-indexing this object exists to prevent."""
    layout = CondLayout.from_types("physical", "physical")
    covered = list(range(*layout.base.indices(COND_DIM)))
    covered += list(range(*layout.particle_phys.indices(COND_DIM)))
    covered += list(range(*layout.material_phys.indices(COND_DIM)))
    assert covered == list(range(COND_DIM))


def test_cont_slice_widths_match_constants():
    layout = CondLayout.from_types("embedding", "embedding")
    assert layout.base == slice(0, COND_DIM_BASE)
    assert layout.particle_phys.stop - layout.particle_phys.start == PARTICLE_PHYS_DIM
    assert layout.material_phys.stop - layout.material_phys.start == MATERIAL_PHYS_DIM
    assert layout.cont_dim == COND_DIM


def test_cont_slices_are_mode_independent():
    """cond_cont is COND_DIM wide in every mode — a non-"physical" axis gets
    its block zero-filled rather than dropped, so the slices never move."""
    physical = CondLayout.from_types("physical", "physical")
    for particle in AXIS_TYPES:
        for material in AXIS_TYPES:
            layout = CondLayout.from_types(particle, material)
            assert layout.base == physical.base
            assert layout.particle_phys == physical.particle_phys
            assert layout.material_phys == physical.material_phys


# ── validation ───────────────────────────────────────────────────────────────


def test_unknown_particle_type_raises():
    with pytest.raises(ValueError, match="unknown conditioning.particle.type 'bogus'"):
        CondLayout.from_types("bogus", "physical")


def test_unknown_material_type_raises():
    with pytest.raises(ValueError, match="unknown conditioning.material.type 'bogus'"):
        CondLayout.from_types("physical", "bogus")