"""Tests for the meta-controller & Lyapunov energy (``palimseste.meta``).

Core guarantees under test:
  - a proposal that reduces surprise is accepted; one that increases it is rejected
  - accepted rewrites are audit-logged into H_meta (append-only)
  - invariants penalize configs (e.g. max_radius blocks runaway radius)
  - the acceptance criterion itself is never rewritten (constitution immutability)
  - build_replay draws from M's own traces
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste import hv
from palimseste.memory import Memory
from palimseste.phi import Phi, KernelConfig
from palimseste.meta import (
    MetaController,
    LyapunovEnergy,
    MetaProposal,
    Invariant,
    max_radius_invariant,
    _config_to_hv,
)


def _setup(D=2000, seed=1, radius=5, n_traces=40):
    rng = np.random.default_rng(seed)
    mem = Memory(D=D, rng=np.random.default_rng(seed))
    # populate M with (address, value) pairs where value == address+noise,
    # so a wider radius genuinely helps reconstruction (lower surprise).
    base = hv.random_hv(D=D, rng=rng)
    for _ in range(n_traces):
        a = _flip(base, 50, rng)
        v = _flip(a, 3, rng)  # value near its address
        mem.write(a, v)
    phi = Phi(config=KernelConfig(radius=radius, min_weight=1e-6))
    energy = LyapunovEnergy()
    ctrl = MetaController(mem=mem, phi=phi, energy=energy, rng=rng)
    return mem, phi, ctrl, rng


def _flip(h: hv.HV, n: int, rng) -> hv.HV:
    s = hv.bits_to_signs(h)
    pos = rng.choice(h.D, size=n, replace=False)
    s[pos] = -s[pos]
    return hv.signs_to_bits(s)


def _noisy_replay(mem, n, flip_bits, rng):
    """Replay set of (noisy_query, target_value): query is a stored address
    with ``flip_bits`` positions flipped, so radius=0 misses but a wider
    radius finds the true value. This is the *generalization* regime where
    the kernel's radius actually matters."""
    idx = rng.choice(len(mem.traces), size=min(n, len(mem.traces)), replace=False)
    replay = []
    for i in idx:
        tr = mem.traces[i]
        q = _flip(tr.address, flip_bits, rng)
        replay.append((q, tr.value))
    return replay


def test_evaluate_accepts_surprise_reducing_proposal():
    mem, phi, ctrl, rng = _setup(radius=0, n_traces=40)
    # noisy queries: 15 bits flipped -> radius 0 misses, radius 20 hits
    replay = _noisy_replay(mem, 40, flip_bits=15, rng=rng)
    prop = MetaProposal(KernelConfig(radius=20, min_weight=1e-6), "widen radius")
    dec = ctrl.evaluate(prop, replay)
    assert dec.accepted
    assert dec.delta < 0


def test_evaluate_rejects_surprise_increasing_proposal():
    mem, phi, ctrl, rng = _setup(radius=30, n_traces=40)
    # with radius 30, noisy queries are found; shrinking to 0 breaks them
    replay = _noisy_replay(mem, 40, flip_bits=15, rng=rng)
    prop = MetaProposal(KernelConfig(radius=0, min_weight=1e-6), "shrink radius")
    dec = ctrl.evaluate(prop, replay)
    assert not dec.accepted
    assert dec.delta > 0


def test_accepted_rewrite_is_audit_logged_to_hmeta():
    mem, phi, ctrl, rng = _setup(radius=0, n_traces=30)
    n_meta_before = len(mem.meta_traces)
    replay = ctrl.build_replay(30)
    prop = MetaProposal(KernelConfig(radius=15, min_weight=1e-6), "widen")
    dec = ctrl.evaluate(prop, replay)
    assert dec.accepted
    # commit manually (evaluate doesn't commit; step does)
    phi.config = dec.proposal.config
    ctrl._write_current_config()
    assert len(mem.meta_traces) == n_meta_before + 1


def test_step_commits_first_accepted_proposal():
    mem, phi, ctrl, rng = _setup(radius=0, n_traces=30)
    replay = ctrl.build_replay(30)
    cfg_before = phi.config
    dec = ctrl.step(replay, max_proposals=20)
    # at least one widening proposal should be accepted eventually
    assert dec is not None
    if dec.accepted:
        assert phi.config == dec.proposal.config
        assert phi.config != cfg_before


def test_invariant_max_radius_blocks_runaway():
    # with a tiny max_radius invariant, widening beyond it is penalized
    mem, phi, ctrl, rng = _setup(radius=0, n_traces=20)
    ctrl.energy = LyapunovEnergy(invariants=[max_radius_invariant(max_r=5)])
    replay = ctrl.build_replay(20)
    # propose radius 100 -> invariant violation huge -> rejected
    prop = MetaProposal(KernelConfig(radius=100, min_weight=1e-6), "runaway")
    dec = ctrl.evaluate(prop, replay)
    assert not dec.accepted
    assert "rejected" in dec.reason


def test_invariant_zero_when_satisfied():
    inv = max_radius_invariant(max_r=50)
    mem = Memory(D=500, rng=np.random.default_rng(0))
    cfg_ok = KernelConfig(radius=10)
    cfg_bad = KernelConfig(radius=100)
    assert inv.violation(mem, cfg_ok) == 0.0
    assert inv.violation(mem, cfg_bad) > 0.0


def test_constitution_immutability():
    # The acceptance criterion (LyapunovEnergy + invariants) is NOT stored in
    # H_meta and cannot be rewritten by the controller. Verify the controller
    # has no API to mutate its own energy/invariants.
    mem, phi, ctrl, rng = _setup(radius=5, n_traces=10)
    # The controller exposes `energy` as a field, but there is no method to
    # *write* a new energy into H_meta. The only thing it writes to H_meta is
    # kernel configs. Assert the meta-traces are all tagged kernel_config:*.
    for tr in mem.meta_traces:
        assert tr.tag is not None and tr.tag.startswith("kernel_config:")


def test_build_replay_from_empty_memory():
    mem = Memory(D=500, rng=np.random.default_rng(0))
    phi = Phi(config=KernelConfig(radius=5))
    ctrl = MetaController(mem=mem, phi=phi, energy=LyapunovEnergy(),
                          rng=np.random.default_rng(0))
    assert ctrl.build_replay(10) == []


def test_config_to_hv_deterministic_and_distinct():
    D = 1000
    rng = np.random.default_rng(0)
    a = KernelConfig(radius=5, min_weight=1e-3, sharpness=0.0, topk=None)
    b = KernelConfig(radius=5, min_weight=1e-3, sharpness=0.0, topk=None)
    c = KernelConfig(radius=6, min_weight=1e-3, sharpness=0.0, topk=None)
    ha = _config_to_hv(a, D, rng)
    hb = _config_to_hv(b, D, rng)
    hc = _config_to_hv(c, D, rng)
    assert ha == hb
    assert ha != hc


def test_history_records_all_evaluated_proposals():
    mem, phi, ctrl, rng = _setup(radius=0, n_traces=15)
    replay = ctrl.build_replay(15)
    ctrl.step(replay, max_proposals=5)
    # at least 1, at most 5 decisions recorded
    assert 1 <= len(ctrl.history) <= 5


def test_custom_invariant_callable():
    called = {"n": 0}

    def _viol(_mem, _cfg):
        called["n"] += 1
        return 0.0

    inv = Invariant(name="custom", violation=_viol, lam=2.0)
    energy = LyapunovEnergy(invariants=[inv])
    mem = Memory(D=500, rng=np.random.default_rng(0))
    phi = Phi(config=KernelConfig(radius=2))
    ctrl = MetaController(mem=mem, phi=phi, energy=energy, rng=np.random.default_rng(0))
    # write a trace so replay is non-empty
    rng = np.random.default_rng(1)
    mem.write(hv.random_hv(D=500, rng=rng), hv.random_hv(D=500, rng=rng))
    replay = ctrl.build_replay(1)
    ctrl.evaluate(MetaProposal(KernelConfig(radius=3), "test"), replay)
    assert called["n"] >= 1