from __future__ import annotations

import json
import random

from demos.grownet_trace_main import run_for_test
from focus import AnchorMap, FocusCandidate, FocusConfig, FocusEngine
from region import Region
from runner.experiment_loader import load_experiment_config_file


def test_scalar_anchor_metadata_is_explicit_after_tick():
    region_obj = Region("scalar_anchor")
    layer_index = region_obj.add_layer(1, 0, 0)
    region_obj.bind_input("in", [layer_index])

    region_obj.tick("in", 0.4)
    region_obj.tick("in", 0.6)

    input_edge_index = region_obj.input_edges["in"]
    neuron_obj = region_obj.get_layers()[input_edge_index].get_neurons()[0]
    assert neuron_obj.last_slot_anchor_state is not None
    assert neuron_obj.last_slot_anchor_state.scalar_value == 0.4
    assert neuron_obj.last_slot_focus_selection is not None
    assert neuron_obj.last_slot_focus_selection.candidate.scalar_value == 0.6
    assert neuron_obj.last_slot_focus_selection.distance_from_anchor > 0.0


def test_focus_engine_generates_sparse_candidates_and_energy_first_selection():
    focus_engine = FocusEngine(
        FocusConfig(enabled=True, policy="energy_first", candidate_limit=4),
        random.Random(1234),
    )
    frame_2d = [
        [0.0, 0.7],
        [0.2, 0.0],
    ]

    focus_selections = focus_engine.plan_frame(frame_2d, tick_number=1)

    assert focus_engine.summarize().candidate_count == 2
    assert focus_selections[0].candidate.point == (0, 1)


def test_random_candidate_policy_is_deterministic_for_same_seed():
    frame_2d = [
        [0.4, 0.8],
        [0.6, 0.2],
    ]
    first_engine = FocusEngine(
        FocusConfig(enabled=True, policy="random_candidate", candidate_limit=1),
        random.Random(4321),
    )
    second_engine = FocusEngine(
        FocusConfig(enabled=True, policy="random_candidate", candidate_limit=1),
        random.Random(4321),
    )

    first_selection = first_engine.plan_frame(frame_2d, tick_number=1)[0]
    second_selection = second_engine.plan_frame(frame_2d, tick_number=1)[0]

    assert first_selection.candidate.point == second_selection.candidate.point


def test_novelty_and_familiarity_policies_respect_anchor_map_distance():
    remembered_candidate = FocusCandidate(point=(0, 0), energy=1.0)

    novelty_engine = FocusEngine(
        FocusConfig(enabled=True, policy="novelty_first", candidate_limit=1),
        random.Random(1234),
    )
    novelty_engine.anchor_map.touch_candidate(remembered_candidate, tick_number=1, policy_name="energy_first")
    novelty_engine.active_anchor_state = novelty_engine.anchor_map.states()[0]

    familiarity_engine = FocusEngine(
        FocusConfig(enabled=True, policy="familiarity_first", candidate_limit=1),
        random.Random(1234),
    )
    familiarity_engine.anchor_map.touch_candidate(remembered_candidate, tick_number=1, policy_name="energy_first")
    familiarity_engine.active_anchor_state = familiarity_engine.anchor_map.states()[0]

    frame_2d = [
        [0.0, 0.4, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0, 0.9],
    ]

    novelty_selection = novelty_engine.plan_frame(frame_2d, tick_number=2)[0]
    familiarity_selection = familiarity_engine.plan_frame(frame_2d, tick_number=2)[0]

    assert novelty_selection.candidate.point == (3, 3)
    assert familiarity_selection.candidate.point == (0, 1)


def test_inhibition_of_return_prefers_other_candidate_on_next_tick():
    focus_engine = FocusEngine(
        FocusConfig(
            enabled=True,
            policy="energy_first",
            candidate_limit=1,
            inhibition_of_return_ticks=3,
        ),
        random.Random(1234),
    )
    frame_2d = [[1.0, 0.6]]

    first_selection = focus_engine.plan_frame(frame_2d, tick_number=1)[0]
    second_selection = focus_engine.plan_frame(frame_2d, tick_number=2)[0]

    assert first_selection.candidate.point == (0, 0)
    assert second_selection.candidate.point == (0, 1)


def test_anchor_map_create_revisit_and_capacity_eviction_are_deterministic():
    anchor_map = AnchorMap(capacity=2)

    first_action, _ = anchor_map.touch_candidate(FocusCandidate(point=(1, 1), energy=0.5), 1, "energy_first")
    second_action, revisited_anchor = anchor_map.touch_candidate(FocusCandidate(point=(1, 1), energy=0.8), 2, "energy_first")
    third_action, _ = anchor_map.touch_candidate(FocusCandidate(point=(2, 2), energy=0.3), 3, "energy_first")
    fourth_action, _ = anchor_map.touch_candidate(FocusCandidate(point=(3, 3), energy=0.4), 4, "energy_first")

    assert first_action == "create"
    assert second_action == "revisit"
    assert revisited_anchor.visit_count == 2
    assert third_action == "create"
    assert fourth_action == "create"
    assert anchor_map.count() == 2
    assert (1, 1) not in [anchor_state.point for anchor_state in anchor_map.states()]


def test_tick_2d_focused_preserves_baseline_when_focus_is_disabled():
    baseline_region = build_spatial_region()
    focused_region = build_spatial_region()
    frame_2d = [
        [0.0, 0.8],
        [0.0, 0.0],
    ]

    baseline_metrics = baseline_region.tick_2d("img", frame_2d)
    focused_metrics = focused_region.tick_2d_focused("img", frame_2d)

    assert focused_metrics.delivered_events == baseline_metrics.delivered_events
    assert focused_metrics.total_slots == baseline_metrics.total_slots
    assert focused_metrics.total_synapses == baseline_metrics.total_synapses


def test_tick_2d_focused_selects_highest_energy_candidate_when_enabled():
    region_obj = build_spatial_region()
    region_obj.set_focus_config(
        FocusConfig(enabled=True, policy="energy_first", serial=False, candidate_limit=1)
    )
    frame_2d = [
        [0.3, 0.0],
        [0.0, 0.9],
    ]

    metrics = region_obj.tick_2d_focused("img", frame_2d)
    focus_summary = region_obj.get_focus_summary()

    assert metrics.delivered_events >= 1
    assert focus_summary.active_focus_point == (1, 1)
    assert focus_summary.anchor_count == 1


def test_focus_trace_events_emit_for_configured_rank_two_experiment(tmp_path):
    config_path = tmp_path / "focus_trace.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "focus_trace_probe",
                "runtime": {"ticks": 1},
                "focus": {
                    "enabled": True,
                    "policy": "energy_first",
                    "candidate_limit": 2,
                    "serial": True,
                    "allow_anchor_map": True,
                    "inhibition_of_return_ticks": 2,
                },
                "input_space": {"rank": 2, "shape": [2, 2], "ports": [{"name": "image"}]},
                "network": {
                    "regions": [
                        {
                            "name": "vision",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ],
                            "input_bindings": {"image": [0]},
                        }
                    ]
                },
                "input": {"type": "tensor", "values": [[0.9, 0.0], [0.0, 0.4]]},
            }
        ),
        encoding="utf-8",
    )

    output_text = run_for_test(
        [
            "--lang",
            "python",
            "--trace",
            "--trace-format",
            "jsonl",
            "--trace-events",
            "focus,anchor",
            "--config",
            str(config_path),
        ]
    )
    payloads = [json.loads(line_value) for line_value in output_text.splitlines() if line_value.strip()]

    assert [payload["event"] for payload in payloads[:3]] == [
        "focus.candidates",
        "focus.select",
        "focus.shift",
    ]
    assert any(payload["event"] == "anchor.create" for payload in payloads)
    assert payloads[0]["candidate_count"] == 2
    assert payloads[1]["focus_policy"] == "energy_first"
    assert payloads[1]["focus_point"] == [0, 0]


def test_experiment_loader_parses_focus_block(tmp_path):
    config_path = tmp_path / "focus_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "focus_config_probe",
                "runtime": {"ticks": 1},
                "focus": {
                    "enabled": True,
                    "policy": "novelty_first",
                    "serial": True,
                    "candidate_limit": 4,
                    "allow_anchor_map": True,
                    "inhibition_of_return_ticks": 5,
                },
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.4]},
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))

    assert experiment_config.focus.enabled is True
    assert experiment_config.focus.policy == "novelty_first"
    assert experiment_config.focus.candidate_limit == 4


def build_spatial_region() -> Region:
    region_obj = Region("focus_spatial")
    input_layer_index = region_obj.add_input_layer_2d(2, 2, 1.0, 0.01)
    hidden_layer_index = region_obj.add_layer(1, 0, 0)
    region_obj.connect_layers(input_layer_index, hidden_layer_index, probability=1.0, feedback=False)
    region_obj.bind_input("img", [input_layer_index])
    return region_obj
