from __future__ import annotations

import numpy as np

from terrarium.map_memory import FIFOLocalMapMemory, MAP_MEMORY_CHANNELS


def _obs(center, terrain, visibility, object_channels=None):
    size = len(terrain)
    return {
        "schema_version": "large_world_local_observation_v1",
        "center": list(center),
        "radius": size // 2,
        "size": size,
        "terrain": terrain,
        "biomes": [["meadow" for _ in range(size)] for _ in range(size)],
        "visibility": visibility,
        "object_channels": object_channels or {},
        "scalar_features": {
            "map_width": 96,
            "map_height": 96,
            "known_world_fraction": size * size / float(96 * 96),
            "visible_fraction": sum(sum(row) for row in visibility) / float(size * size),
        },
    }


def test_fifo_memory_only_stores_visible_cells():
    observation = _obs(
        (10, 10),
        [
            ["wall", "open", "open"],
            ["open", "open", "open"],
            ["open", "open", "hazard"],
        ],
        [
            [0.0, 1.0, 1.0],
            [1.0, 1.0, 1.0],
            [1.0, 1.0, 0.0],
        ],
    )
    memory = FIFOLocalMapMemory(capacity=50)
    memory.update(observation, step=1)

    remembered = memory.remembered_cells()
    assert (9, 9) not in remembered
    assert (11, 11) not in remembered
    assert (10, 10) in remembered
    assert remembered[(10, 10)].visited is True


def test_fifo_memory_evicts_oldest_percepts_deterministically():
    memory = FIFOLocalMapMemory(capacity=2)
    for step, center in enumerate(((0, 0), (1, 0), (2, 0))):
        memory.update(_obs(center, [["open"]], [[1.0]]), step=step)

    remembered = memory.remembered_cells()
    assert memory.oldest_step == 1
    assert memory.newest_step == 2
    assert (0, 0) not in remembered
    assert {(1, 0), (2, 0)} <= set(remembered)


def test_fifo_memory_local_channels_are_egocentric_and_resettable():
    source = _obs(
        (10, 10),
        [["shelter"]],
        [[1.0]],
        {"shelter": [[1.0]], "edible": [[0.0]], "hazard": [[0.0]], "novelty": [[0.25]]},
    )
    current = _obs(
        (11, 10),
        [
            ["open", "open", "open"],
            ["open", "open", "open"],
            ["open", "open", "open"],
        ],
        [
            [1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0],
        ],
    )
    memory = FIFOLocalMapMemory(capacity=50)
    memory.update(source, step=1)

    channels = memory.local_channels(current)
    assert channels.shape == (len(MAP_MEMORY_CHANNELS), 3, 3)
    shelter_index = MAP_MEMORY_CHANNELS.index("remembered_shelter")
    novelty_index = MAP_MEMORY_CHANNELS.index("remembered_novelty")
    # Remembered cell (10, 10) is one step west of current center (11, 10).
    assert channels[shelter_index, 1, 0] > 0.0
    assert channels[novelty_index, 1, 0] > 0.0
    assert channels.dtype == np.float32

    memory.clear()
    assert len(memory) == 0
    assert memory.local_channels(current).sum() == 0.0