package ai.nektron.grownet.tests;

import ai.nektron.grownet.Layer;
import ai.nektron.grownet.LateralBus;
import ai.nektron.grownet.Neuron;
import ai.nektron.grownet.SlotConfig;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;

public class LayerForwardSnapshotRegressionTests {

    @Test
    void growthDuringForwardDoesNotReenterNewNeuronUntilNextTick() {
        ForwardSnapshotOutcome firstOutcome = runForwardSnapshotScenario();
        ForwardSnapshotOutcome secondOutcome = runForwardSnapshotScenario();

        assertEquals(firstOutcome, secondOutcome);
        assertEquals(2, firstOutcome.layerNeuronCountAfterFirstTick);
        assertEquals(0, firstOutcome.grownNeuronInputsAfterFirstTick);
        assertEquals(1, firstOutcome.grownNeuronInputsAfterSecondTick);
    }

    private static ForwardSnapshotOutcome runForwardSnapshotScenario() {
        SnapshotAwareLayer layer = new SnapshotAwareLayer();
        GrowingSeedNeuron seedNeuron = new GrowingSeedNeuron("seed", layer.getBus(), SlotConfig.fixed(10.0), -1);
        seedNeuron.owner = layer;
        layer.getNeurons().add(seedNeuron);

        assertDoesNotThrow(() -> layer.forward(1.0));
        CountingNeuron grownNeuron = layer.getGrownNeuron();
        assertNotNull(grownNeuron);
        int layerNeuronCountAfterFirstTick = layer.getNeurons().size();
        int grownNeuronInputsAfterFirstTick = grownNeuron.getInputCount();

        assertDoesNotThrow(() -> layer.forward(1.0));
        int grownNeuronInputsAfterSecondTick = grownNeuron.getInputCount();
        return new ForwardSnapshotOutcome(
                layerNeuronCountAfterFirstTick,
                grownNeuronInputsAfterFirstTick,
                grownNeuronInputsAfterSecondTick
        );
    }

    private static final class SnapshotAwareLayer extends Layer {
        private CountingNeuron grownNeuron;

        private SnapshotAwareLayer() {
            super(0, 0, 0);
        }

        @Override
        public int tryGrowNeuron(Neuron seedNeuron) {
            int newNeuronIndex = getNeurons().size();
            grownNeuron = new CountingNeuron(
                    "grown_" + newNeuronIndex,
                    getBus(),
                    seedNeuron.getSlotConfig().copy(),
                    seedNeuron.getSlotLimit()
            );
            grownNeuron.owner = this;
            getNeurons().add(grownNeuron);
            return newNeuronIndex;
        }

        private CountingNeuron getGrownNeuron() {
            return grownNeuron;
        }
    }

    private static class CountingNeuron extends Neuron {
        private int inputCount = 0;

        private CountingNeuron(String neuronId, LateralBus bus, SlotConfig slotConfig, int slotLimit) {
            super(neuronId, bus, slotConfig, slotLimit);
        }

        @Override
        public boolean onInput(double value) {
            inputCount += 1;
            return false;
        }

        private int getInputCount() {
            return inputCount;
        }
    }

    private static final class GrowingSeedNeuron extends CountingNeuron {
        private boolean alreadyGrown = false;

        private GrowingSeedNeuron(String neuronId, LateralBus bus, SlotConfig slotConfig, int slotLimit) {
            super(neuronId, bus, slotConfig, slotLimit);
        }

        @Override
        public boolean onInput(double value) {
            super.onInput(value);
            if (!alreadyGrown && owner != null) {
                alreadyGrown = true;
                owner.tryGrowNeuron(this);
            }
            return false;
        }
    }

    private static final class ForwardSnapshotOutcome {
        private final int layerNeuronCountAfterFirstTick;
        private final int grownNeuronInputsAfterFirstTick;
        private final int grownNeuronInputsAfterSecondTick;

        private ForwardSnapshotOutcome(
                int layerNeuronCountAfterFirstTick,
                int grownNeuronInputsAfterFirstTick,
                int grownNeuronInputsAfterSecondTick) {
            this.layerNeuronCountAfterFirstTick = layerNeuronCountAfterFirstTick;
            this.grownNeuronInputsAfterFirstTick = grownNeuronInputsAfterFirstTick;
            this.grownNeuronInputsAfterSecondTick = grownNeuronInputsAfterSecondTick;
        }

        @Override
        public boolean equals(Object otherObject) {
            if (this == otherObject) {
                return true;
            }
            if (!(otherObject instanceof ForwardSnapshotOutcome otherOutcome)) {
                return false;
            }
            return layerNeuronCountAfterFirstTick == otherOutcome.layerNeuronCountAfterFirstTick
                    && grownNeuronInputsAfterFirstTick == otherOutcome.grownNeuronInputsAfterFirstTick
                    && grownNeuronInputsAfterSecondTick == otherOutcome.grownNeuronInputsAfterSecondTick;
        }

        @Override
        public int hashCode() {
            return java.util.Objects.hash(
                    layerNeuronCountAfterFirstTick,
                    grownNeuronInputsAfterFirstTick,
                    grownNeuronInputsAfterSecondTick
            );
        }
    }
}
