package ai.nektron.grownet;

import ai.nektron.grownet.debug.Trace;

import java.util.LinkedHashMap;
import java.util.Map;

public final class SlotEngine {
    private final SlotConfig cfg;

    public SlotEngine(SlotConfig cfg) { this.cfg = cfg; }

    public SlotConfig getConfig() { return cfg; }

    private static int chooseExistingSlotId(Neuron neuron, int preferredSlotId) {
        if (neuron.getSlots().containsKey(preferredSlotId)) {
            return preferredSlotId;
        }
        int selectedSlotId = preferredSlotId;
        boolean foundExistingSlot = false;
        for (Integer existingSlotId : neuron.getSlots().keySet()) {
            if (existingSlotId == null) continue;
            if (!foundExistingSlot || existingSlotId < selectedSlotId) {
                selectedSlotId = existingSlotId;
                foundExistingSlot = true;
            }
        }
        return selectedSlotId;
    }

    public int slotId(double lastInput, double newInput, int slotsLen) {
        double deltaPercent = 0.0;
        if (lastInput != 0.0) {
            deltaPercent = Math.abs(newInput - lastInput) / Math.abs(lastInput) * 100.0;
        }
        switch (cfg.policy) {
            case FIXED:      return fixedBucket(deltaPercent);
            case NONUNIFORM: return nonuniformBucket(deltaPercent);
            case ADAPTIVE:   return fixedBucket(deltaPercent); // base id; overflow handled by caller
        }
        return 0;
    }

    private int fixedBucket(double deltaPercent) {
        if (deltaPercent <= 0.0) return 0;
        return (int)Math.floor((deltaPercent + (cfg.slotWidthPercent - 1.0)) / cfg.slotWidthPercent);
    }

    private int nonuniformBucket(double deltaPercent) {
        int idx = 0;
        for (double edge : cfg.nonuniformEdges) {
            if (deltaPercent <= edge) return idx;
            idx++;
        }
        return idx;
    }

    /**
     * Temporal‑focus helper (FIRST anchor): choose a slot id for input x, ensure a slot exists,
     * and clamp growth at cfg.getSlotLimit() when at capacity.
     */
    /** FIRST-anchor helper with strict capacity + fallback marking.
     *  Chooses desired bin; if out-of-domain or at capacity and new, falls back to a deterministic id.
     *  Never allocates a brand-new slot at capacity (unless this is the very first slot).
     */
    public int selectOrCreateSlot(Neuron neuron, double inputValue, SlotConfig cfg) {
        if (neuron == null) return 0;
        if (cfg == null) cfg = this.cfg;

        // Anchor (FIRST)
        if (!neuron.focusSet && cfg.getAnchorMode() == SlotConfig.AnchorMode.FIRST) {
            neuron.focusAnchor = inputValue;
            neuron.focusSet = true;
        }

        double anchor = neuron.focusAnchor;
        double denom = Math.max(Math.abs(anchor), cfg.getEpsilonScale());
        double deltaPct = Math.abs(inputValue - anchor) / denom * 100.0;
        double bin = Math.max(0.1, cfg.getBinWidthPct());
        int sidDesired = (int) Math.floor(deltaPct / bin);

        int effectiveLimit = (neuron.slotLimit >= 0 ? neuron.slotLimit : cfg.getSlotLimit());
        boolean atCapacity = (effectiveLimit > 0 && neuron.getSlots().size() >= effectiveLimit);
        boolean outOfDomain = (effectiveLimit > 0 && sidDesired >= effectiveLimit);
        boolean wantNew = !neuron.getSlots().containsKey(sidDesired);
        boolean useFallback = outOfDomain || (atCapacity && wantNew);
        int slotsBefore = neuron.getSlots().size();

        int sid = useFallback && effectiveLimit > 0 ? (effectiveLimit - 1) : sidDesired;
        if (useFallback && atCapacity && !neuron.getSlots().isEmpty()) {
            sid = chooseExistingSlotId(neuron, sid);
        }
        boolean createdSlot = false;
        if (!neuron.getSlots().containsKey(sid)) {
            if (atCapacity) {
                if (neuron.getSlots().isEmpty()) {
                    neuron.getSlots().put(sid, new Weight());
                    createdSlot = true;
                }
            } else {
                neuron.getSlots().put(sid, new Weight());
                createdSlot = true;
            }
        }
        neuron.lastSlotUsedFallback = useFallback;
        if (useFallback) {
            neuron.lastMissingSlotId = sidDesired;
            neuron.lastMaxAxisDeltaPct = deltaPct;
        } else {
            neuron.fallbackStreak = 0;
            neuron.prevMissingSlotId = -1;
            neuron.lastMissingSlotId = -1;
        }
        Map<String, Object> selectionEvent = new LinkedHashMap<>();
        Trace.addNeuronFields(selectionEvent, neuron);
        selectionEvent.put("desired_slot", sidDesired);
        selectionEvent.put("chosen_slot", sid);
        selectionEvent.put("fallback", useFallback);
        selectionEvent.put("created_slot", createdSlot);
        selectionEvent.put("slots_before", slotsBefore);
        selectionEvent.put("slots_after", neuron.getSlots().size());
        selectionEvent.put("slot_limit", effectiveLimit);
        Trace.emit("slot.select", selectionEvent);
        if (createdSlot) {
            Map<String, Object> createdEvent = new LinkedHashMap<>();
            Trace.addNeuronFields(createdEvent, neuron);
            createdEvent.put("slot_id", sid);
            createdEvent.put("created_slot", true);
            Trace.emit("slot.create", createdEvent);
        }
        return sid;
    }

    /**
     * Spatial (2D) helper with strict capacity + fallback marking.
     * Uses FIRST spatial anchor; packs (rowBin,colBin) as rowBin*100000 + colBin.
     */
    public int selectOrCreateSlot2D(Neuron neuron, int row, int col, SlotConfig cfg) {
        if (neuron == null) return 0;
        if (cfg == null) cfg = this.cfg;

        // FIRST spatial anchor
        if (neuron.anchorRow < 0 || neuron.anchorCol < 0) {
            neuron.anchorRow = row;
            neuron.anchorCol = col;
        }

        double rowDenom = Math.max(Math.abs(neuron.anchorRow), Math.max(1.0, cfg.getEpsilonScale()));
        double colDenom = Math.max(Math.abs(neuron.anchorCol), Math.max(1.0, cfg.getEpsilonScale()));
        double rowDeltaPct = Math.abs(row - neuron.anchorRow) / rowDenom * 100.0;
        double colDeltaPct = Math.abs(col - neuron.anchorCol) / colDenom * 100.0;
        int rowBin = (int) Math.floor(rowDeltaPct / Math.max(0.1, cfg.getRowBinWidthPct()));
        int colBin = (int) Math.floor(colDeltaPct / Math.max(0.1, cfg.getColBinWidthPct()));

        int effectiveLimit = (neuron.slotLimit >= 0 ? neuron.slotLimit : cfg.getSlotLimit());
        boolean atCapacity = (effectiveLimit > 0 && neuron.getSlots().size() >= effectiveLimit);
        boolean outOfDomain = (effectiveLimit > 0 && (rowBin >= effectiveLimit || colBin >= effectiveLimit));

        int desiredKey = rowBin * 100000 + colBin;
        boolean wantNew = !neuron.getSlots().containsKey(desiredKey);
        boolean useFallback = outOfDomain || (atCapacity && wantNew);
        int slotsBefore = neuron.getSlots().size();

        int key = useFallback && effectiveLimit > 0
                ? (effectiveLimit - 1) * 100000 + (effectiveLimit - 1)
                : desiredKey;
        if (useFallback && atCapacity && !neuron.getSlots().isEmpty()) {
            key = chooseExistingSlotId(neuron, key);
        }

        boolean createdSlot = false;
        if (!neuron.getSlots().containsKey(key)) {
            if (atCapacity) {
                if (neuron.getSlots().isEmpty()) {
                    neuron.getSlots().put(key, new Weight());
                    createdSlot = true;
                }
            } else {
                neuron.getSlots().put(key, new Weight());
                createdSlot = true;
            }
        }
        neuron.lastSlotUsedFallback = useFallback;
        if (useFallback) {
            neuron.lastMissingSlotId = desiredKey;
            neuron.lastMaxAxisDeltaPct = Math.max(rowDeltaPct, colDeltaPct);
        } else {
            neuron.fallbackStreak = 0;
            neuron.prevMissingSlotId = -1;
            neuron.lastMissingSlotId = -1;
        }
        Map<String, Object> selectionEvent = new LinkedHashMap<>();
        Trace.addNeuronFields(selectionEvent, neuron);
        selectionEvent.put("desired_slot", desiredKey);
        selectionEvent.put("chosen_slot", key);
        selectionEvent.put("fallback", useFallback);
        selectionEvent.put("created_slot", createdSlot);
        selectionEvent.put("slots_before", slotsBefore);
        selectionEvent.put("slots_after", neuron.getSlots().size());
        selectionEvent.put("slot_limit", effectiveLimit);
        Trace.emit("slot.select", selectionEvent);
        if (createdSlot) {
            Map<String, Object> createdEvent = new LinkedHashMap<>();
            Trace.addNeuronFields(createdEvent, neuron);
            createdEvent.put("slot_id", key);
            createdEvent.put("created_slot", true);
            Trace.emit("slot.create", createdEvent);
        }
        return key;
    }
}
