package ai.nektron.grownet.tests;

import ai.nektron.grownet.ExcitatoryNeuron;
import ai.nektron.grownet.InhibitoryNeuron;
import ai.nektron.grownet.InputLayer2D;
import ai.nektron.grownet.InputLayerND;
import ai.nektron.grownet.LateralBus;
import ai.nektron.grownet.Layer;
import ai.nektron.grownet.Neuron;
import ai.nektron.grownet.Region;
import ai.nektron.grownet.SlotConfig;
import ai.nektron.grownet.SlotEngine;
import ai.nektron.grownet.Weight;
import ai.nektron.grownet.growth.GrowthPolicy;
import org.junit.jupiter.api.Test;

import java.util.List;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertTrue;

final class ParityStabilizationTests {

    private static int totalOutgoingSynapses(Layer layer) {
        int totalOutgoingSynapses = 0;
        for (Neuron neuron : layer.getNeurons()) {
            totalOutgoingSynapses += neuron.getOutgoing().size();
        }
        return totalOutgoingSynapses;
    }

    @Test
    void saturatedScalarFallbackReusesAnExistingSlotId() {
        LateralBus layerBus = new LateralBus();
        SlotConfig slotConfig = SlotConfig.fixed(10.0).setSlotLimit(2);
        ExcitatoryNeuron sourceNeuron = new ExcitatoryNeuron("scalar-seed", layerBus, slotConfig, 2);
        SlotEngine slotEngine = new SlotEngine(slotConfig);

        slotEngine.selectOrCreateSlot(sourceNeuron, 10.0, slotConfig);
        sourceNeuron.getSlots().put(7, new Weight());

        int selectedSlotId = slotEngine.selectOrCreateSlot(sourceNeuron, 50.0, slotConfig);

        assertEquals(0, selectedSlotId, "fallback should reuse the smallest existing slot id");
        assertTrue(sourceNeuron.getSlots().containsKey(selectedSlotId), "returned slot must already exist");
        assertEquals(2, sourceNeuron.getSlots().size(), "capacity saturation must not allocate a third slot");
        assertTrue(sourceNeuron.lastSlotUsedFallback, "saturated reuse must still be marked as fallback");
    }

    @Test
    void saturatedSpatialFallbackReusesAnExistingPackedSlotId() {
        LateralBus layerBus = new LateralBus();
        SlotConfig slotConfig = SlotConfig.fixed(10.0)
                .setSlotLimit(2)
                .setRowBinWidthPct(10.0)
                .setColBinWidthPct(10.0);
        ExcitatoryNeuron sourceNeuron = new ExcitatoryNeuron("spatial-seed", layerBus, slotConfig, 2);
        SlotEngine slotEngine = new SlotEngine(slotConfig);

        slotEngine.selectOrCreateSlot2D(sourceNeuron, 10, 10, slotConfig);
        sourceNeuron.getSlots().put(700007, new Weight());

        int selectedSlotId = slotEngine.selectOrCreateSlot2D(sourceNeuron, 50, 50, slotConfig);

        assertEquals(0, selectedSlotId, "2D fallback should reuse the smallest existing packed slot id");
        assertTrue(sourceNeuron.getSlots().containsKey(selectedSlotId), "returned 2D slot must already exist");
        assertEquals(2, sourceNeuron.getSlots().size(), "2D saturation must not allocate a third slot");
        assertTrue(sourceNeuron.lastSlotUsedFallback, "2D saturated reuse must still be marked as fallback");
    }

    @Test
    void grownInhibitoryNeuronCopiesSeedConfigurationAndSubtypeState() {
        Layer layer = new Layer(0, 1, 0);
        InhibitoryNeuron seedNeuron = assertInstanceOf(InhibitoryNeuron.class, layer.getNeurons().get(0));
        seedNeuron.getSlotConfig()
                .setBinWidthPct(12.5)
                .setRowBinWidthPct(25.0)
                .setColBinWidthPct(40.0)
                .setFallbackGrowthThreshold(5)
                .setLayerGrowthEnabled(true);
        seedNeuron.setGamma(0.33);

        int grownNeuronIndex = layer.tryGrowNeuron(seedNeuron);
        InhibitoryNeuron grownNeuron = assertInstanceOf(
                InhibitoryNeuron.class,
                layer.getNeurons().get(grownNeuronIndex)
        );

        assertNotSame(seedNeuron.getSlotConfig(), grownNeuron.getSlotConfig(), "growth must copy slot config");
        assertEquals(12.5, grownNeuron.getSlotConfig().getBinWidthPct(), 1e-12);
        assertEquals(25.0, grownNeuron.getSlotConfig().getRowBinWidthPct(), 1e-12);
        assertEquals(40.0, grownNeuron.getSlotConfig().getColBinWidthPct(), 1e-12);
        assertEquals(5, grownNeuron.getSlotConfig().getFallbackGrowthThreshold());
        assertTrue(grownNeuron.getSlotConfig().isLayerGrowthEnabled());
        assertEquals(0.33, grownNeuron.getGamma(), 1e-12, "inhibitory subtype parameters must be preserved");
    }

    @Test
    void bindInput2dSkipsSelfEdgeConnectionsWhenRebinding() {
        Region region = new Region("two-dimensional-self-edge");
        int targetLayerIndex = region.addLayer(1, 0, 0);

        region.bindInput2D("image", 2, 2, 1.0, 0.01, List.of(targetLayerIndex));
        int edgeLayerIndex = region.getLayers().size() - 1;
        Layer edgeLayer = assertInstanceOf(InputLayer2D.class, region.getLayers().get(edgeLayerIndex));
        int outgoingBeforeRebind = totalOutgoingSynapses(edgeLayer);

        region.bindInput2D("image", 2, 2, 1.0, 0.01, List.of(edgeLayerIndex));

        assertEquals(outgoingBeforeRebind, totalOutgoingSynapses(edgeLayer));
    }

    @Test
    void bindInputNdSkipsSelfEdgeConnectionsWhenRebinding() {
        Region region = new Region("tensor-self-edge");
        int targetLayerIndex = region.addLayer(1, 0, 0);

        region.bindInputND("tensor", new int[]{2, 2}, 1.0, 0.01, List.of(targetLayerIndex));
        int edgeLayerIndex = region.getLayers().size() - 1;
        Layer edgeLayer = assertInstanceOf(InputLayerND.class, region.getLayers().get(edgeLayerIndex));
        int outgoingBeforeRebind = totalOutgoingSynapses(edgeLayer);

        region.bindInputND("tensor", new int[]{2, 2}, 1.0, 0.01, List.of(edgeLayerIndex));

        assertEquals(outgoingBeforeRebind, totalOutgoingSynapses(edgeLayer));
    }

    @Test
    void growthPolicyDefaultsRemainConservativeAndAligned() {
        GrowthPolicy growthPolicy = new GrowthPolicy();

        assertFalse(growthPolicy.isEnableLayerGrowth());
        assertEquals(16, growthPolicy.getMaxLayers());
        assertEquals(8.0, growthPolicy.getAvgSlotsThreshold(), 1e-12);
        assertEquals(50.0, growthPolicy.getPercentAtCapFallbackThreshold(), 1e-12);
        assertEquals(25L, growthPolicy.getLayerCooldownTicks());
        assertEquals(4, growthPolicy.getNewLayerExcitatoryCount());
        assertEquals(1.0, growthPolicy.getWireProbability(), 1e-12);
    }
}
