from focus.focus_engine import add_active_focus_fields, update_scalar_slot_anchor_metadata, update_spatial_slot_anchor_metadata
from slot_config import SlotPolicy, normalize_anchor_mode
from weight import Weight
from typing import Tuple
from debug import trace as grownet_trace


def pack_slot_key_2d(slot_key: Tuple[int, int] | None) -> int | None:
    if slot_key is None:
        return None
    return int(int(slot_key[0]) * 100000 + int(slot_key[1]))


def slot_key_to_trace_id(slot_key) -> int | None:
    if slot_key is None:
        return None
    if isinstance(slot_key, tuple):
        return pack_slot_key_2d(slot_key)
    return int(slot_key)


class SlotEngine:
    """Slot selection helpers (policy + temporal & spatial focus) with fallback markers."""
    def __init__(self, cfg):
        self.cfg = cfg

    # -------- scalar (temporal) --------
    def slot_id(self, last_input, current_input, known_slots):
        """Map a percent delta to an integer bin (policy-dependent)."""
        if self.cfg.policy == SlotPolicy.FIXED:
            if last_input == 0:
                return 0
            delta_percent = abs(current_input - last_input) / max(1e-9, abs(last_input)) * 100.0
            bin_width = max(1.0, self.cfg.fixed_delta_percent)
            return int(delta_percent // bin_width)
        elif self.cfg.policy == SlotPolicy.NONUNIFORM:
            if last_input == 0:
                return 0
            delta_percent = abs(current_input - last_input) / max(1e-9, abs(last_input)) * 100.0
            for index, edge in enumerate(self.cfg.custom_edges):
                if delta_percent < edge:
                    return index
            return len(self.cfg.custom_edges)
        else:
            # ADAPTIVE or unknown: collapse everything to slot 0
            return 0

    def select_or_create_slot(self, neuron, input_value, tick_count=0):
        """FIRST-anchor binning with capacity clamp.
        Marks `neuron.last_slot_used_fallback` whenever a new bin is desired but capacity forces reuse.
        Never allocates a brand-new slot if capacity is already reached.
        """
        cfg = self.cfg
        anchor_mode = normalize_anchor_mode(getattr(cfg, "anchor_mode", "FIRST"), spatial_enabled=False)
        # FIRST-anchor: set anchor once
        if not getattr(neuron, "focus_set", False) and anchor_mode == "FIRST":
            neuron.focus_anchor = float(input_value)
            neuron.focus_set = True

        anchor = float(getattr(neuron, "focus_anchor", 0.0))
        eps = max(1e-12, float(getattr(cfg, "epsilon_scale", 1e-6)))
        denom = max(abs(anchor), eps)
        delta_pct = abs(float(input_value) - anchor) / denom * 100.0
        bin_w = max(0.1, float(getattr(cfg, "bin_width_pct", 10.0)))
        sid_desired = int(delta_pct // bin_w)
        update_scalar_slot_anchor_metadata(neuron, float(input_value), anchor, float(delta_pct), anchor_mode)

        # Effective limit: per-neuron override wins over cfg
        neuron_limit = int(getattr(neuron, "slot_limit", -1))
        limit = neuron_limit if neuron_limit >= 0 else int(getattr(cfg, "slot_limit", 16))
        slots = neuron.slots
        at_capacity = (limit > 0 and len(slots) >= limit)
        slots_before = len(slots)

        want_new = (sid_desired not in slots)
        out_of_domain = (limit > 0 and sid_desired >= limit)
        use_fallback = out_of_domain or (at_capacity and want_new)

        # Choose actual sid
        if use_fallback and limit > 0:
            sid = (limit - 1)
            if sid not in slots and slots:
                sid = next(iter(slots.keys()))
        else:
            sid = sid_desired

        # Ensure existence (respect capacity)
        created_slot = False
        if sid not in slots:
            if at_capacity:
                if not slots:
                    slots[sid] = Weight()
                    created_slot = True
            else:
                slots[sid] = Weight()
                created_slot = True

        try:
            neuron.last_selected_slot_id = sid
            neuron.last_slot_used_fallback = bool(use_fallback)
            if use_fallback:
                neuron.last_missing_slot_id = sid_desired
                neuron.last_max_axis_delta_pct = float(delta_pct)
            else:
                neuron.fallback_streak = 0
                neuron.prev_missing_slot_id = None
                neuron.last_missing_slot_id = None
        except Exception:
            pass
        selection_event = {}
        grownet_trace.add_neuron_fields(selection_event, neuron)
        selection_event["desired_slot"] = sid_desired
        selection_event["chosen_slot"] = sid
        selection_event["fallback"] = bool(use_fallback)
        selection_event["created_slot"] = bool(created_slot)
        selection_event["slots_before"] = int(slots_before)
        selection_event["slots_after"] = len(slots)
        selection_event["slot_limit"] = int(limit)
        add_active_focus_fields(selection_event, neuron)
        grownet_trace.emit("slot.select", **selection_event)
        if created_slot:
            created_event = {}
            grownet_trace.add_neuron_fields(created_event, neuron)
            created_event["slot_id"] = sid
            created_event["created_slot"] = True
            add_active_focus_fields(created_event, neuron)
            grownet_trace.emit("slot.create", **created_event)
        return slots[sid]

    # -------- spatial (2D) --------
    def slot_id_2d(self, anchor_rc: Tuple[int, int], current_rc: Tuple[int, int]) -> Tuple[int, int]:
        """Return (row_bin, col_bin) using |Δrow| and |Δcol| as % of denom (per-axis).

        Denominator per axis uses max(|anchor_axis|, epsilon_scale) to avoid divide-by-zero.
        Bin widths are controlled by cfg.row_bin_width_pct / cfg.col_bin_width_pct.
        """
        anchor_row, anchor_col = int(anchor_rc[0]), int(anchor_rc[1])
        current_row, current_col = int(current_rc[0]), int(current_rc[1])

        # Use a sensible spatial epsilon to avoid exploding bins at ORIGIN (0,0).
        eps = max(1.0, float(getattr(self.cfg, "epsilon_scale", 1.0)))
        denominator_row = max(abs(anchor_row), eps)
        denominator_col = max(abs(anchor_col), eps)

        delta_percent_row = abs(current_row - anchor_row) / denominator_row * 100.0
        delta_percent_col = abs(current_col - anchor_col) / denominator_col * 100.0

        bin_width_row_pct = max(0.1, float(getattr(self.cfg, "row_bin_width_pct", 100.0)))
        bin_width_col_pct = max(0.1, float(getattr(self.cfg, "col_bin_width_pct", 100.0)))

        return int(delta_percent_row // bin_width_row_pct), int(delta_percent_col // bin_width_col_pct)

    def select_or_create_slot_2d(self, neuron, row: int, col: int):
        """2D FIRST/ORIGIN anchor + capacity clamp; ensures spatial slot exists.

        Keys are (row_bin, col_bin) tuples. When capacity is saturated, reuse a
        fallback id (limit-1, limit-1) if present; otherwise reuse any existing slot.
        Never creates a new slot at capacity. Also sets neuron.last_slot_used_fallback.
        """
        cfg = self.cfg

        # pick anchors
        anchor_mode = normalize_anchor_mode(getattr(cfg, "anchor_mode", "FIRST"), spatial_enabled=True)
        if anchor_mode == "ORIGIN":
            anchor_row, anchor_col = 0, 0
        else:
            # FIRST: set if not present
            if getattr(neuron, "focus_anchor_row", None) is None or getattr(neuron, "focus_anchor_col", None) is None:
                neuron.focus_anchor_row = int(row)
                neuron.focus_anchor_col = int(col)
            anchor_row, anchor_col = int(neuron.focus_anchor_row), int(neuron.focus_anchor_col)

        row_bin, col_bin = self.slot_id_2d((anchor_row, anchor_col), (int(row), int(col)))
        limit = int(getattr(cfg, "slot_limit", 16))

        if limit > 0:
            row_bin = min(row_bin, limit - 1)
            col_bin = min(col_bin, limit - 1)

        key_desired = (row_bin, col_bin)
        slots = neuron.slots
        # Effective limit: per-neuron override wins over cfg
        neuron_limit = int(getattr(neuron, "slot_limit", -1))
        limit = neuron_limit if neuron_limit >= 0 else int(getattr(cfg, "slot_limit", 16))
        at_capacity = (limit > 0 and len(slots) >= limit)
        slots_before = len(slots)
        want_new = (key_desired not in slots)
        use_fallback = (at_capacity and want_new)

        if use_fallback and limit > 0:
            key = (limit - 1, limit - 1)
            if key not in slots:
                if slots:
                    key = next(iter(slots.keys()))
        else:
            key = key_desired

        created_slot = False
        if key not in slots:
            if at_capacity:
                if not slots:
                    slots[key] = Weight()
                    created_slot = True
            else:
                slots[key] = Weight()
                created_slot = True
        try:
            neuron.last_selected_slot_id = key
            neuron.last_slot_used_fallback = bool(use_fallback)
            if use_fallback:
                eps = max(1.0, float(getattr(cfg, "epsilon_scale", 1.0)))
                row_delta_pct = abs(int(row) - anchor_row) / max(abs(anchor_row), eps) * 100.0
                col_delta_pct = abs(int(col) - anchor_col) / max(abs(anchor_col), eps) * 100.0
                neuron.last_max_axis_delta_pct = max(row_delta_pct, col_delta_pct)
                neuron.last_missing_slot_id = key_desired
            else:
                neuron.fallback_streak = 0
                neuron.prev_missing_slot_id = None
                neuron.last_missing_slot_id = None
        except Exception:
            pass
        update_spatial_slot_anchor_metadata(
            neuron,
            int(row),
            int(col),
            (int(anchor_row), int(anchor_col)),
            float(max(abs(int(row) - anchor_row), abs(int(col) - anchor_col))),
            anchor_mode,
        )
        selection_event = {}
        grownet_trace.add_neuron_fields(selection_event, neuron)
        selection_event["desired_slot"] = pack_slot_key_2d(key_desired)
        selection_event["chosen_slot"] = pack_slot_key_2d(key)
        selection_event["fallback"] = bool(use_fallback)
        selection_event["created_slot"] = bool(created_slot)
        selection_event["slots_before"] = int(slots_before)
        selection_event["slots_after"] = len(slots)
        selection_event["slot_limit"] = int(limit)
        add_active_focus_fields(selection_event, neuron)
        grownet_trace.emit("slot.select", **selection_event)
        if created_slot:
            created_event = {}
            grownet_trace.add_neuron_fields(created_event, neuron)
            created_event["slot_id"] = pack_slot_key_2d(key)
            created_event["created_slot"] = True
            add_active_focus_fields(created_event, neuron)
            grownet_trace.emit("slot.create", **created_event)
        return slots[key]
