import pytest
from region import Region


def get_first_neuron(region: Region, layer_index: int):
    """Helper to read the first neuron of a layer."""
    return region.get_layers()[layer_index].get_neurons()[0]


def last_slot_of(neuron):
    """
    Robustly fetch the most recently used slot object.
    Falls back to any slot if the convenience pointer is absent.
    """
    slot = getattr(neuron, "last_slot", None)
    if slot is not None:
        return slot
    # Fallback (shouldn't be needed if PR for frozen slots is applied):
    slots_map = getattr(neuron, "slots", {})
    assert isinstance(slots_map, dict) and slots_map, "Neuron has no slots yet"
    return next(iter(slots_map.values()))


def test_frozen_slot_scalar_stops_adaptation_and_unfreeze_resumes():
    """
    Drive a scalar edge into a 1-neuron hidden layer.
    - After the first tick we freeze the last-selected slot.
    - While frozen, reinforcement and theta must NOT change.
    - After unfreezing, reinforcement resumes (strength increases).
    """
    region = Region("frozen_scalar")
    target_layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    region.bind_input("x", [target_layer_index])
    edge_layer_index = region.input_edges["x"]

    # First tick establishes FIRST anchor + picks/creates a slot.
    region.tick("x", 0.6)
    source_neuron = get_first_neuron(region, edge_layer_index)
    selected_slot = last_slot_of(source_neuron)

    # Record baseline before freezing.
    initial_strength = float(selected_slot.strength)
    initial_theta = float(selected_slot.theta)

    # Freeze and tick again (values would normally adapt).
    assert source_neuron.freeze_last_slot() is True
    region.tick("x", 0.6)

    # While frozen: no reinforcement, no theta drift.
    assert float(selected_slot.strength) == pytest.approx(initial_strength)
    assert float(selected_slot.theta) == pytest.approx(initial_theta)

    # Unfreeze and tick again — strength should now grow by the step (~0.02).
    assert source_neuron.unfreeze_last_slot() is True
    region.tick("x", 0.6)
    assert float(selected_slot.strength) > initial_strength  # adaptation resumed


def test_frozen_slot_spatial_2d_stops_adaptation_and_unfreeze_resumes():
    """
    Same idea as the scalar test, but through the spatial (row,col) path.
    We enable spatial slotting on a single hidden neuron and deliver the
    same pixel twice while frozen — strength/theta must not change.
    """
    region = Region("frozen_2d")
    input_layer_index = region.add_input_layer_2d(3, 3, gain=1.0, epsilon_fire=0.01)
    hidden_layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)

    # Enable spatial slotting on the destination neuron.
    source_neuron = get_first_neuron(region, hidden_layer_index)
    source_neuron.slot_cfg.spatial_enabled = True
    source_neuron.slot_cfg.row_bin_width_pct = 100.0
    source_neuron.slot_cfg.col_bin_width_pct = 100.0

    # Deterministic 2D wiring and bind input.
    region.connect_layers_windowed(input_layer_index, hidden_layer_index, kernel_h=3, kernel_w=3, stride_h=3, stride_w=3, padding="valid")
    region.bind_input("img", [input_layer_index])

    # Drive a single bright pixel at the same location to hit the same spatial bin.
    frame = [[0.0 for _ in range(3)] for _ in range(3)]
    frame[1][1] = 1.0

    region.tick_2d("img", frame)  # establish spatial slot
    selected_slot = last_slot_of(source_neuron)
    initial_strength = float(selected_slot.strength)
    initial_theta = float(selected_slot.theta)

    # Freeze -> tick (same pixel) -> no changes.
    assert source_neuron.freeze_last_slot() is True
    region.tick_2d("img", frame)
    assert float(selected_slot.strength) == pytest.approx(initial_strength)
    assert float(selected_slot.theta) == pytest.approx(initial_theta)

    # Unfreeze -> tick (same pixel) -> reinforcement resumes.
    assert source_neuron.unfreeze_last_slot() is True
    region.tick_2d("img", frame)
    assert float(selected_slot.strength) > initial_strength
