"""Trace geometry, segmentation and persistence, without a model."""
import numpy as np
import pytest

import ulens
from ulens import runtime
from ulens.runtime import Trace, _segment

OPEN, CLOSE = 900, 901


def make_trace(n_blocks=5, T=20, k=4, prompt_len=6, eot=14, seed=0, **kw):
    rng = np.random.default_rng(seed)
    opts: dict = {
        "model": "qwen3.5-4b", "prompt": "p",
        "tokens": list(range(T)), "token_strs": [f"t{i}" for i in range(T)],
        "prompt_len": prompt_len, "eot": eot,
        "entropy": np.concatenate(
            [[np.nan], rng.uniform(0.1, 2.0, size=T - 1)]).astype(np.float32),
        "field": rng.normal(size=(n_blocks, T, k)).astype(np.float32),
        "blocks": list(range(10, 10 + n_blocks)),
        "names": ulens.terms.categories()[:k],
        "readout_block": 12, "generated": "g", "truncated": False,
        "config": {"seed": 1}}
    opts.update(kw)
    return Trace(**opts)


# ------------------------------------------------------------------ segmentation

def test_block_already_open_at_the_prompt_end():
    """The usual case: the chat template opens the channel, the generation closes it."""
    ids = [1, 2, OPEN] + [5, 6, 7, CLOSE, 8, 9]
    eot, trunc = _segment(ids, 3, OPEN, CLOSE)
    assert eot == 6 and ids[eot] == CLOSE and not trunc


def test_block_opened_inside_the_generation():
    ids = [1, 2, 3] + [OPEN, 5, 6, CLOSE, 8]
    eot, trunc = _segment(ids, 3, OPEN, CLOSE)
    assert eot == 6 and ids[eot] == CLOSE and not trunc


def test_opened_and_never_closed_is_all_thinking():
    """What a token budget produces. It must not read as a zero-token answer."""
    ids = [1, 2, OPEN] + [5, 6, 7, 8]
    eot, trunc = _segment(ids, 3, OPEN, CLOSE)
    assert eot == len(ids) - 1 and trunc
    tr = make_trace(T=len(ids), prompt_len=3, eot=eot, truncated=trunc)
    assert tr.think_len == len(ids) - 3
    assert (tr.segments == "answer").sum() == 0


def test_no_thinking_block_at_all_is_all_answer():
    ids = [1, 2, 3] + [5, 6, 7]
    eot, trunc = _segment(ids, 3, OPEN, CLOSE)
    assert eot == 2 and not trunc
    tr = make_trace(T=6, prompt_len=3, eot=eot)
    assert tr.think_len == 0
    assert list(tr.segments) == ["prompt"] * 3 + ["answer"] * 3


def test_a_closer_in_the_prompt_does_not_count():
    """A prompt that mentions the closer must not end the block before it starts."""
    ids = [1, CLOSE, 2, OPEN] + [5, 6, CLOSE, 9]
    eot, _ = _segment(ids, 4, OPEN, CLOSE)
    assert eot == 6


def test_a_reopened_block_uses_the_last_opener():
    ids = [OPEN, 1, CLOSE, 2, OPEN] + [5, 6, CLOSE]
    eot, trunc = _segment(ids, 5, OPEN, CLOSE)
    assert eot == 7 and not trunc


def test_without_a_closer_token_everything_is_thinking():
    ids = [1, 2, 3, 4, 5]
    eot, trunc = _segment(ids, 2, OPEN, None)
    assert eot == len(ids) - 1 and trunc


# ------------------------------------------------------------------ geometry

def test_segments_partition_the_sequence():
    tr = make_trace()
    s = tr.segments
    assert len(s) == tr.n_tokens
    assert (s == "prompt").sum() == tr.prompt_len
    assert (s == "think").sum() == tr.think_len
    assert (s == "answer").sum() == tr.n_tokens - tr.eot - 1


def test_a_cone_series_matches_the_readout_at_the_eot():
    tr = make_trace()
    assert tr.a_cone == pytest.approx(tr.a_cone_series()[tr.eot])
    z = tr.projections()[tr.eot]
    assert tr.a_cone == pytest.approx(float(ulens.a_cone_of(z)))


def test_the_field_agrees_with_a_cone_computed_from_states():
    """a_cone_of on stored projections equals a_cone on the states they came from."""
    rng = np.random.default_rng(1)
    U = ulens._math.lowdin(rng.normal(size=(32, 4)).astype(np.float32))
    H = rng.normal(size=(7, 32)).astype(np.float32)
    z = ulens.projections(H, U)
    np.testing.assert_allclose(ulens.a_cone_of(z), ulens.a_cone(H, U), rtol=1e-6)


def test_surface_is_a_cone_over_depth_and_token():
    tr = make_trace()
    S = tr.surface()
    assert S.shape == (len(tr.blocks), tr.n_tokens)
    np.testing.assert_allclose(S[tr.block_index()], tr.a_cone_series(), rtol=1e-6)


def test_components_are_the_rectified_projections():
    tr = make_trace()
    np.testing.assert_allclose(tr.components(), ulens.rectify(tr.projections()), rtol=1e-6)
    assert np.all(tr.components() > 0)


def test_think_entropy_spans_the_thinking_block_inclusive_of_the_eot():
    tr = make_trace()
    want = tr.entropy[tr.prompt_len: tr.eot + 1].mean()
    assert tr.think_entropy == pytest.approx(want)
    assert tr.u_lens == pytest.approx(tr.a_cone * tr.think_entropy)


def test_the_first_token_has_no_entropy_and_it_is_excluded():
    """Nothing predicted token 0, so its entropy is NaN — and must not reach the mean."""
    tr = make_trace()
    assert np.isnan(tr.entropy[0])
    assert not np.isnan(tr.think_entropy)


def test_an_uncaptured_block_is_refused():
    tr = make_trace()
    with pytest.raises(KeyError, match="was not captured"):
        tr.projections(999)


def test_summary_and_axis_table_render():
    tr = make_trace()
    assert "U-Lens" in tr.summary() and "readout block 12" in tr.summary()
    assert "closed" in tr.summary() and tr.readout_valid
    assert "argmax" in tr.axis_table()


def test_a_truncated_trace_says_the_readout_is_not_the_published_one():
    tr = make_trace(truncated=True)
    assert not tr.readout_valid
    assert "TRUNCATED" in tr.summary() and "not the published one" in tr.summary()
    assert np.isfinite(tr.u_lens)          # still computed; the field is real


def test_a_trace_with_no_thinking_block_is_not_a_valid_readout():
    tr = make_trace(prompt_len=6, eot=5)
    assert tr.think_len == 0 and not tr.readout_valid


# ------------------------------------------------------------------ persistence

def test_trace_round_trips_and_needs_no_model(tmp_path):
    tr = make_trace()
    back = runtime.load(tr.save(tmp_path / "t.npz"))
    np.testing.assert_allclose(back.field, tr.field)
    np.testing.assert_array_equal(np.isnan(back.entropy), np.isnan(tr.entropy))
    assert back.u_lens == pytest.approx(tr.u_lens)
    assert back.tokens == tr.tokens and back.token_strs == tr.token_strs
    assert back.blocks == tr.blocks and back.names == tr.names
    assert back.eot == tr.eot and back.truncated == tr.truncated
    assert back.config == tr.config


def test_the_sidecar_carries_the_scores(tmp_path):
    import json
    tr = make_trace()
    p = tr.save(tmp_path / "t.npz")
    side = json.loads(p.with_suffix(".json").read_text())
    assert side["scores"]["u_lens"] == pytest.approx(tr.u_lens)
    assert side["readout_block"] == tr.readout_block


def test_the_sidecar_is_strict_json_without_a_thinking_block(tmp_path):
    """Thinking off: no thinking-block entropy, no score. The sidecar says null, not NaN."""
    import json
    tr = make_trace(prompt_len=6, eot=5)
    assert tr.think_len == 0
    p = tr.save(tmp_path / "t.npz")
    text = p.with_suffix(".json").read_text()
    assert "NaN" not in text
    side = json.loads(text, parse_constant=lambda c: pytest.fail(f"non-finite {c} in sidecar"))
    assert side["scores"]["think_entropy"] is None and side["scores"]["u_lens"] is None
    assert side["scores"]["a_cone"] == pytest.approx(tr.a_cone)


# ------------------------------------------------------------------ cone stacks

def test_cone_stack_round_trips_and_slices(tmp_path):
    rng = np.random.default_rng(2)
    cats = ulens.terms.categories()
    cones = {b: ulens.axes.from_axes(
        [ulens.Axis(c, rng.normal(size=12).astype(np.float32)) for c in cats],
        model="m", block=b) for b in (3, 5, 9)}
    st = ulens.axes.stack(cones, model="m", meta={"terms_version": "v1"})
    assert st.blocks == [3, 5, 9] and st.U.shape == (3, 12, 4)
    assert st.nearest(6) == 5 and st.nearest(100) == 9
    np.testing.assert_array_equal(st.at(5).U, cones[5].U)
    back = ulens.axes.load_stack(st.save(tmp_path / "s.npz"))
    np.testing.assert_allclose(back.U, st.U)
    assert back.blocks == st.blocks and back.names == st.names


def test_cone_stack_rejects_mismatched_axis_names():
    v = np.ones(8, np.float32)
    a = ulens.axes.from_axes([ulens.Axis("x", v)], block=1)
    b = ulens.axes.from_axes([ulens.Axis("y", v)], block=2)
    with pytest.raises(ValueError, match="expected"):
        ulens.axes.stack({1: a, 2: b})
    with pytest.raises(ValueError, match="at least one cone"):
        ulens.axes.stack({})


def test_stack_index_refuses_an_uncovered_block():
    v = np.ones(8, np.float32)
    st = ulens.axes.stack({4: ulens.axes.from_axes([ulens.Axis("x", v)], block=4)})
    with pytest.raises(KeyError, match="not in this stack"):
        st.index(5)