import pytest
import time
from neo_cortex.surprisal import surprisal_score, select_by_surprisal


def test_surprisal_high_energy_low_access():
    """High energy + low access count = high surprisal."""
    score = surprisal_score(energy=0.9, access_count=0, days_old=3)
    assert score > 0.2


def test_surprisal_low_energy():
    """Low energy = low surprisal regardless of access."""
    score = surprisal_score(energy=0.1, access_count=0, days_old=3)
    assert score < 0.1


def test_surprisal_high_access_reduces_score():
    """Frequently accessed memories have lower surprisal."""
    low_access = surprisal_score(energy=0.8, access_count=1, days_old=5)
    high_access = surprisal_score(energy=0.8, access_count=20, days_old=5)
    assert low_access > high_access


def test_surprisal_very_old_reduces_score():
    """Very old memories have lower surprisal (not novel anymore)."""
    recent = surprisal_score(energy=0.8, access_count=2, days_old=1)
    old = surprisal_score(energy=0.8, access_count=2, days_old=100)
    assert recent > old


def test_surprisal_zero_energy():
    """Zero energy = zero surprisal."""
    score = surprisal_score(energy=0.0, access_count=0, days_old=1)
    assert score == 0.0


def test_surprisal_score_range():
    """Surprisal score is always in [0, 1]."""
    assert 0.0 <= surprisal_score(energy=1.0, access_count=0, days_old=0) <= 1.0
    assert 0.0 <= surprisal_score(energy=0.0, access_count=100, days_old=1000) <= 1.0
    assert 0.0 <= surprisal_score(energy=0.5, access_count=5, days_old=30) <= 1.0


def test_select_by_surprisal_returns_top_n():
    """select_by_surprisal returns top N memory IDs by score."""
    memories = [
        {"id": "a", "energy": 0.9, "access_count": 0, "timestamp": time.time() - 86400},
        {"id": "b", "energy": 0.3, "access_count": 10, "timestamp": time.time() - 86400 * 30},
        {"id": "c", "energy": 0.7, "access_count": 1, "timestamp": time.time() - 86400 * 2},
    ]
    result = select_by_surprisal(memories, n=2)
    assert len(result) == 2
    # "a" should rank highest (high energy, zero access, recent)
    assert result[0] == "a"


def test_select_by_surprisal_empty():
    """select_by_surprisal with empty list returns empty."""
    assert select_by_surprisal([], n=3) == []


def test_select_by_surprisal_respects_n():
    """select_by_surprisal returns at most n results."""
    memories = [
        {"id": f"m{i}", "energy": 0.5, "access_count": 0, "timestamp": time.time()}
        for i in range(10)
    ]
    result = select_by_surprisal(memories, n=3)
    assert len(result) == 3
