from __future__ import annotations

import random

from debug import trace as grownet_trace

from .anchor_map import AnchorMap
from .focus_policy import normalize_focus_policy_name, select_candidate
from .focus_types import AnchorState, FocusCandidate, FocusConfig, FocusSelection, FocusTickSummary, point_payload


class FocusEngine:
    def __init__(self, focus_config: FocusConfig | None = None, rng: random.Random | None = None) -> None:
        self.config = focus_config or FocusConfig()
        self.rng = rng or random.Random(1234)
        self.anchor_map = AnchorMap(capacity=int(self.config.anchor_map_capacity))
        self.active_anchor_state: AnchorState | None = None
        self.active_selection: FocusSelection | None = None
        self.last_summary = FocusTickSummary()
        self.last_focus_ticks: dict[tuple[int, int], int] = {}

    def reset(self) -> None:
        self.anchor_map.clear()
        self.active_anchor_state = None
        self.active_selection = None
        self.last_summary = FocusTickSummary()
        self.last_focus_ticks = {}

    def is_enabled(self) -> bool:
        return bool(self.config.enabled)

    def anchor_count(self) -> int:
        return self.anchor_map.count() if bool(self.config.allow_anchor_map) else 0

    def summarize(self) -> FocusTickSummary:
        return self.last_summary

    def plan_frame(self, frame_2d, tick_number: int) -> list[FocusSelection]:
        if not self.is_enabled():
            self.last_summary = FocusTickSummary()
            return []

        self.anchor_map.expire_before(int(tick_number) - int(self.config.anchor_decay_ticks))
        candidates = self.generate_candidates(frame_2d, tick_number)
        self._emit_candidates_event(candidates)
        if not candidates:
            self.active_selection = None
            self.last_summary = FocusTickSummary(
                policy_name=self.config.policy,
                candidate_count=0,
                selection_count=0,
                anchor_count=self.anchor_count(),
            )
            return []

        serial_limit = max(1, int(self.config.candidate_limit))
        if not bool(self.config.serial):
            serial_limit = 1

        remaining_candidates = list(candidates)
        selections: list[FocusSelection] = []
        previous_focus_point = self.active_selection.candidate.point if self.active_selection is not None else None
        for selected_rank in range(1, min(serial_limit, len(remaining_candidates)) + 1):
            available_candidates, inhibited_candidates = self._split_inhibited_candidates(remaining_candidates, tick_number)
            if available_candidates:
                candidate_pool = available_candidates
            else:
                candidate_pool = remaining_candidates

            if inhibited_candidates and available_candidates:
                for inhibited_candidate in inhibited_candidates:
                    self._emit_return_inhibited(inhibited_candidate)

            selected_candidate = select_candidate(candidate_pool, self.config.policy, self.rng)
            anchor_before = self.active_anchor_state
            anchor_after, anchor_event_name = self._update_anchor_state(selected_candidate, tick_number)
            distance_from_anchor = _distance_from_anchor(anchor_before, selected_candidate)
            selection = FocusSelection(
                candidate=selected_candidate,
                policy_name=self.config.policy,
                selected_rank=int(selected_rank),
                anchor_before=anchor_before,
                anchor_after=anchor_after,
                distance_from_anchor=distance_from_anchor,
                inhibited=selected_candidate not in available_candidates and bool(available_candidates),
            )
            self.active_selection = selection
            self.active_anchor_state = anchor_after
            self.last_focus_ticks[selected_candidate.point] = int(tick_number) if selected_candidate.point is not None else int(tick_number)
            self._emit_selection_events(selection, previous_focus_point, anchor_event_name)
            previous_focus_point = selected_candidate.point
            selections.append(selection)
            remaining_candidates.remove(selected_candidate)

        selected_points = tuple(
            selection.candidate.point for selection in selections if selection.candidate.point is not None
        )
        selected_values = tuple(float(selection.candidate.energy) for selection in selections)
        final_anchor = self.active_anchor_state
        self.last_summary = FocusTickSummary(
            policy_name=self.config.policy,
            candidate_count=len(candidates),
            selection_count=len(selections),
            active_focus_point=selections[-1].candidate.point if selections and selections[-1].candidate.point is not None else None,
            active_focus_value=selections[-1].candidate.scalar_value if selections else None,
            active_anchor_point=final_anchor.point if final_anchor is not None else None,
            active_anchor_value=final_anchor.scalar_value if final_anchor is not None else None,
            anchor_count=self.anchor_count(),
            selected_points=selected_points,
            selected_values=selected_values,
        )
        return selections

    def generate_candidates(self, frame_2d, tick_number: int) -> list[FocusCandidate]:
        candidates: list[FocusCandidate] = []
        for row_index, row_values in enumerate(frame_2d):
            for col_index, pixel_value in enumerate(row_values):
                energy_value = abs(float(pixel_value))
                if energy_value <= 0.0:
                    continue
                candidate_point = (int(row_index), int(col_index))
                novelty_score, familiarity_score = self._scores_for_point(candidate_point, energy_value)
                last_visited_penalty = self._last_visited_penalty(candidate_point, tick_number)
                candidates.append(
                    FocusCandidate(
                        point=candidate_point,
                        energy=energy_value,
                        novelty_score=novelty_score,
                        familiarity_score=familiarity_score,
                        last_visited_penalty=last_visited_penalty,
                    )
                )
        return candidates

    def _scores_for_point(self, point_value: tuple[int, int], energy_value: float) -> tuple[float, float]:
        nearest_anchor = self.active_anchor_state
        nearest_distance = float("inf")
        remembered_anchor, remembered_distance = self.anchor_map.nearest_anchor(point_value)
        if remembered_anchor is not None and remembered_distance < nearest_distance:
            nearest_anchor = remembered_anchor
            nearest_distance = remembered_distance
        if nearest_anchor is None:
            return float(energy_value), 0.0
        if nearest_distance == float("inf"):
            return float(energy_value), 0.0
        novelty_score = float(nearest_distance)
        familiarity_score = 1.0 / (1.0 + float(nearest_distance))
        return novelty_score, familiarity_score

    def _split_inhibited_candidates(
        self,
        candidates: list[FocusCandidate],
        tick_number: int,
    ) -> tuple[list[FocusCandidate], list[FocusCandidate]]:
        available_candidates: list[FocusCandidate] = []
        inhibited_candidates: list[FocusCandidate] = []
        for candidate in candidates:
            if float(candidate.last_visited_penalty) > 0.0:
                inhibited_candidates.append(candidate)
            else:
                available_candidates.append(candidate)
        return available_candidates, inhibited_candidates

    def _last_visited_penalty(self, point_value: tuple[int, int], tick_number: int) -> float:
        if int(self.config.inhibition_of_return_ticks) <= 0:
            return 0.0
        last_tick = self.last_focus_ticks.get(point_value)
        if last_tick is None:
            return 0.0
        elapsed_ticks = int(tick_number) - int(last_tick)
        if elapsed_ticks > int(self.config.inhibition_of_return_ticks):
            return 0.0
        remaining_ticks = int(self.config.inhibition_of_return_ticks) - max(0, elapsed_ticks) + 1
        return float(remaining_ticks)

    def _update_anchor_state(
        self,
        candidate: FocusCandidate,
        tick_number: int,
    ) -> tuple[AnchorState, str]:
        if not bool(self.config.allow_anchor_map):
            anchor_state = AnchorState(
                kind="spatial" if candidate.point is not None else "scalar",
                point=candidate.point,
                scalar_value=candidate.scalar_value,
                created_tick=int(tick_number),
                updated_tick=int(tick_number),
                visit_count=1,
                energy=float(candidate.energy),
                policy_name=self.config.policy,
            )
            return anchor_state, "anchor.update"

        anchor_action, anchor_state = self.anchor_map.touch_candidate(candidate, tick_number, self.config.policy)
        if anchor_action == "create":
            return anchor_state, "anchor.create"
        if anchor_action == "revisit":
            return anchor_state, "anchor.revisit"
        return anchor_state, "anchor.update"

    def _emit_candidates_event(self, candidates: list[FocusCandidate]) -> None:
        if not candidates:
            grownet_trace.emit(
                "focus.candidates",
                focus_policy=self.config.policy,
                candidate_count=0,
            )
            return
        top_candidate = max(
            candidates,
            key=lambda candidate: (
                float(candidate.energy),
                -(candidate.point[0] if candidate.point is not None else 0),
                -(candidate.point[1] if candidate.point is not None else 0),
            ),
        )
        grownet_trace.emit(
            "focus.candidates",
            focus_policy=self.config.policy,
            candidate_count=len(candidates),
            focus_point=point_payload(top_candidate.point),
            candidate_energy=float(top_candidate.energy),
            candidate_novelty=float(top_candidate.novelty_score),
            candidate_familiarity=float(top_candidate.familiarity_score),
        )

    def _emit_selection_events(
        self,
        selection: FocusSelection,
        previous_focus_point: tuple[int, int] | None,
        anchor_event_name: str,
    ) -> None:
        selection_fields = selection_event_fields(selection)
        grownet_trace.emit("focus.select", **selection_fields)
        if previous_focus_point != selection.candidate.point:
            grownet_trace.emit(
                "focus.shift",
                focus_policy=selection.policy_name,
                focus_point=point_payload(selection.candidate.point),
                previous_focus_point=point_payload(previous_focus_point),
                selected_rank=int(selection.selected_rank),
            )
        anchor_fields = selection_event_fields(selection)
        grownet_trace.emit(anchor_event_name, **anchor_fields)

    def _emit_return_inhibited(self, candidate: FocusCandidate) -> None:
        grownet_trace.emit(
            "focus.return_inhibited",
            focus_policy=self.config.policy,
            focus_point=point_payload(candidate.point),
            candidate_energy=float(candidate.energy),
            inhibited=True,
        )


def update_scalar_slot_anchor_metadata(
    neuron,
    input_value: float,
    anchor_value: float,
    distance_from_anchor: float,
    anchor_mode: str,
) -> None:
    anchor_state = AnchorState(
        kind="scalar",
        scalar_value=float(anchor_value),
        created_tick=_current_step_from_neuron(neuron),
        updated_tick=_current_step_from_neuron(neuron),
        visit_count=1 if bool(getattr(neuron, "focus_set", False)) else 0,
        energy=abs(float(input_value)),
        policy_name=str(anchor_mode).lower(),
    )
    selection = FocusSelection(
        candidate=FocusCandidate(
            scalar_value=float(input_value),
            energy=abs(float(input_value)),
        ),
        policy_name="scalar_anchor",
        selected_rank=1,
        anchor_before=anchor_state,
        anchor_after=anchor_state,
        distance_from_anchor=float(distance_from_anchor),
        inhibited=False,
    )
    neuron.last_slot_anchor_state = anchor_state
    neuron.last_slot_focus_selection = selection


def update_spatial_slot_anchor_metadata(
    neuron,
    row_value: int,
    col_value: int,
    anchor_point: tuple[int, int],
    distance_from_anchor: float,
    anchor_mode: str,
) -> None:
    anchor_state = AnchorState(
        kind="spatial",
        point=(int(anchor_point[0]), int(anchor_point[1])),
        created_tick=_current_step_from_neuron(neuron),
        updated_tick=_current_step_from_neuron(neuron),
        visit_count=1,
        energy=0.0,
        policy_name=str(anchor_mode).lower(),
    )
    selection = FocusSelection(
        candidate=FocusCandidate(
            point=(int(row_value), int(col_value)),
            energy=1.0,
        ),
        policy_name="spatial_anchor",
        selected_rank=1,
        anchor_before=anchor_state,
        anchor_after=anchor_state,
        distance_from_anchor=float(distance_from_anchor),
        inhibited=False,
    )
    neuron.last_slot_anchor_state = anchor_state
    neuron.last_slot_focus_selection = selection


def add_active_focus_fields(event_fields: dict, neuron) -> None:
    region_obj = getattr(getattr(neuron, "owner", None), "region", None)
    selection = getattr(region_obj, "active_focus_selection", None)
    if selection is None:
        return
    event_fields.update(selection_event_fields(selection))


def selection_event_fields(selection: FocusSelection) -> dict:
    fields = {
        "focus_policy": selection.policy_name,
        "selected_rank": int(selection.selected_rank),
        "distance_from_anchor": float(selection.distance_from_anchor),
        "candidate_energy": float(selection.candidate.energy),
        "candidate_novelty": float(selection.candidate.novelty_score),
        "candidate_familiarity": float(selection.candidate.familiarity_score),
        "inhibited": bool(selection.inhibited),
    }
    if selection.candidate.point is not None:
        fields["focus_point"] = point_payload(selection.candidate.point)
    if selection.candidate.scalar_value is not None:
        fields["focus_value"] = float(selection.candidate.scalar_value)
    if selection.anchor_before is not None:
        if selection.anchor_before.point is not None:
            fields["anchor_before"] = point_payload(selection.anchor_before.point)
        if selection.anchor_before.scalar_value is not None:
            fields["anchor_before_value"] = float(selection.anchor_before.scalar_value)
    if selection.anchor_after is not None:
        if selection.anchor_after.point is not None:
            fields["anchor_after"] = point_payload(selection.anchor_after.point)
        if selection.anchor_after.scalar_value is not None:
            fields["anchor_after_value"] = float(selection.anchor_after.scalar_value)
    return fields


def _distance_from_anchor(anchor_state: AnchorState | None, candidate: FocusCandidate) -> float:
    if anchor_state is None:
        return 0.0
    if anchor_state.point is not None and candidate.point is not None:
        return float(abs(anchor_state.point[0] - candidate.point[0]) + abs(anchor_state.point[1] - candidate.point[1]))
    if anchor_state.scalar_value is not None and candidate.scalar_value is not None:
        return abs(float(anchor_state.scalar_value) - float(candidate.scalar_value))
    return 0.0


def _current_step_from_neuron(neuron) -> int:
    try:
        bus_obj = neuron.get_bus() if hasattr(neuron, "get_bus") else None
        if bus_obj is not None and hasattr(bus_obj, "get_current_step"):
            return int(bus_obj.get_current_step())
    except Exception:
        return 0
    return 0
