tests/test_loop.py

"""Tests for the autonomous loop (``palimseste.loop``).

Uses a tiny deterministic environment: a 4-state ring where each state's
observation HV is near the next state's (so the loop can learn the transition).
Verifies:
  - the loop runs without error for many ticks
  - surprise decreases over time as the agent learns the transition
  - traces accumulate in M (append-only growth)
  - StepReport telemetry is populated correctly
  - curiosity picks an action (non-None) when actions are available
  - reset_state clears recurrent state but not M
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste import hv
from palimseste.loop import (
    Palimseste,
    Environment,
    StateProjector,
    LoopConfig,
    StepReport,
)
from palimseste.learner import Encoder


# ----------------------------------------------------------------- test env
class RingEnv(Environment):
    """A ring of N states; each tick advances to the next state.

    Observations are HVs that are *near* their neighbors (a few bits apart),
    so the transition o_t -> o_{t+1} is learnable by the associative memory.
    Actions are [advance] (deterministic) — kept minimal so the curiosity
    machinery is exercised without complicating the dynamics.
    """

    def __init__(self, D: int, n_states: int = 4, seed: int = 0):
        self.D = D
        self.n_states = n_states
        rng = np.random.default_rng(seed)
        base = hv.random_hv(D=D, rng=rng)
        self.states: list[hv.HV] = [base]
        signs = hv.bits_to_signs(base)
        bps = max(1, D // 50)
        cur = signs.copy()
        for _ in range(n_states - 1):
            flip = rng.choice(D, size=bps, replace=False)
            cur = cur.copy()
            cur[flip] = -cur[flip]
            self.states.append(hv.signs_to_bits(cur))
        self._i = 0
        self._advance = hv.random_hv(D=D, rng=rng)
        self._t = 0
        self._max_t = 10_000

    def observe(self) -> hv.HV:
        return self.states[self._i]

    def actions(self) -> list[hv.HV]:
        return [self._advance]

    def act(self, action: hv.HV) -> None:
        # single action: advance the ring
        self._i = (self._i + 1) % self.n_states
        self._t += 1

    def done(self) -> bool:
        return self._t >= self._max_t


# ----------------------------------------------------------------- tests
def _agent(D=1500, seed=0, **kw) -> Palimseste:
    return Palimseste(
        D=D,
        rng=np.random.default_rng(seed),
        loop_cfg=LoopConfig(
            surprise_threshold=0.25,
            consolidate_every=16,
            meta_every=64,
            max_radius=80,
        ),
        **kw,
    )


def test_loop_runs_many_ticks():
    agent = _agent(D=1200, seed=1)
    env = RingEnv(D=1200, n_states=4, seed=1)
    reports = []
    for _ in range(200):
        reports.append(agent.step(env))
    assert len(reports) == 200
    assert all(isinstance(r, StepReport) for r in reports)
    # memory grew (append-only learning)
    assert agent.stats()["n_traces"] > 0


def test_surprise_decreases_over_time():
    # After enough ticks the agent should predict the ring transition well,
    # so mean surprise in the second half < mean surprise in the first half.
    agent = _agent(D=1500, seed=2)
    env = RingEnv(D=1500, n_states=4, seed=2)
    surprises = []
    for _ in range(400):
        r = agent.step(env)
        surprises.append(r.surprise)
    first = np.mean(surprises[:100])
    second = np.mean(surprises[300:])
    assert second < first, f"surprise did not decrease: {first=} {second=}"


def test_step_report_fields_populated():
    agent = _agent(D=1000, seed=3)
    env = RingEnv(D=1000, n_states=3, seed=3)
    r = agent.step(env)
    assert r.t == 1
    assert 0.0 <= r.surprise <= 1.0
    assert r.action_idx is not None  # at least one action
    assert r.n_traces >= 0
    assert r.n_concepts == 0  # nothing consolidated yet on tick 1


def test_curiosity_picks_action():
    agent = _agent(D=1000, seed=4)
    env = RingEnv(D=1000, n_states=3, seed=4)
    r = agent.step(env)
    assert r.action_idx == 0  # only one action available


def test_consolidation_fires():
    # With a small consolidate_every, consolidation should run at least once
    agent = _agent(D=1200, seed=5)
    agent.loop_cfg.consolidate_every = 8
    env = RingEnv(D=1200, n_states=4, seed=5)
    ran_cons = False
    for _ in range(40):
        r = agent.step(env)
        if r.consolidation is not None:
            ran_cons = True
    assert ran_cons


def test_meta_fires():
    agent = _agent(D=1200, seed=6)
    agent.loop_cfg.meta_every = 16
    env = RingEnv(D=1200, n_states=4, seed=6)
    ran_meta = False
    for _ in range(80):
        r = agent.step(env)
        if r.meta is not None:
            ran_meta = True
    assert ran_meta


def test_reset_state_clears_recurrence_not_memory():
    agent = _agent(D=1000, seed=7)
    env = RingEnv(D=1000, n_states=3, seed=7)
    for _ in range(30):
        agent.step(env)
    n_before = agent.stats()["n_traces"]
    agent.reset_state()
    n_after = agent.stats()["n_traces"]
    assert n_before == n_after  # M untouched
    assert agent.surprise == 0.0  # reset clears last surprise


def test_loop_config_invalid():
    with pytest.raises(ValueError):
        LoopConfig(surprise_threshold=1.5)
    with pytest.raises(ValueError):
        LoopConfig(consolidate_every=0)
    with pytest.raises(ValueError):
        LoopConfig(max_radius=0)


def test_stats_keys():
    agent = _agent(D=800, seed=8)
    env = RingEnv(D=800, n_states=3, seed=8)
    for _ in range(10):
        agent.step(env)
    s = agent.stats()
    for k in ("t", "n_traces", "n_meta", "n_concepts", "n_meta_decisions",
              "last_surprise", "kernel"):
        assert k in s


def test_state_projector_reset():
    enc = Encoder(D=500, rng=np.random.default_rng(0))
    sp = StateProjector(D=500, encoder=enc, window=3)
    o = hv.random_hv(D=500)
    a = hv.random_hv(D=500)
    # project (state from history, empty at first), then commit o
    s_empty = sp.project(o, a)
    sp.commit(o)
    s_one = sp.project(o, a)  # now history has o
    sp.reset()
    s_after = sp.project(o, a)  # history cleared again
    # s_empty (no history) and s_one (one item in history) differ
    assert hv.similarity(s_empty, s_one) < 1.0
    # after reset, state matches the empty-history state
    assert hv.similarity(s_empty, s_after) == pytest.approx(1.0)