package ai.nektron.grownet.debug;

import ai.nektron.grownet.Layer;
import ai.nektron.grownet.Neuron;
import ai.nektron.grownet.Region;
import ai.nektron.grownet.SlotConfig;
import ai.nektron.grownet.Weight;
import ai.nektron.grownet.growth.GrowthEngine;
import ai.nektron.grownet.growth.GrowthPolicy;
import ai.nektron.grownet.metrics.RegionMetrics;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.io.Writer;
import java.lang.reflect.Field;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.LinkedHashMap;
import java.util.Locale;
import java.util.Map;

public final class GrowNetTraceMain {
    private GrowNetTraceMain() {
    }

    public static void main(String[] args) throws Exception {
        TraceConfig traceConfig = TraceConfig.parse("java", args);
        if (traceConfig.isHelpRequested()) {
            System.out.println(TraceConfig.usage("java"));
            return;
        }
        if (traceConfig.isBenchmarkEnabled()) {
            Writer benchmarkWriter = openBenchmarkWriter(traceConfig);
            Writer traceWriter = openBenchmarkTraceWriter(traceConfig);
            try {
                BenchmarkRunner.run(traceConfig, benchmarkWriter, traceWriter);
            } finally {
                closeWriter(traceConfig.writesBenchmarkToStdout(), benchmarkWriter);
                closeWriter(traceConfig.writesToStdout(), traceWriter);
            }
            return;
        }

        Writer traceWriter = openTraceWriter(traceConfig);
        try {
            run(traceConfig, traceWriter);
        } finally {
            closeWriter(traceConfig.writesToStdout(), traceWriter);
        }
    }

    public static String runForTest(String... args) throws Exception {
        TraceConfig traceConfig = TraceConfig.parse("java", args);
        StringWriter outputBuffer = new StringWriter();
        if (traceConfig.isBenchmarkEnabled()) {
            BenchmarkRunner.run(traceConfig, outputBuffer, new DiscardingWriter());
        } else {
            execute(traceConfig, outputBuffer, true);
        }
        return outputBuffer.toString();
    }

    public static TraceRunSummary runSummaryForTest(String... args) throws Exception {
        return runSummaryForTestWithWriter(new StringWriter(), args);
    }

    public static TraceRunSummary runSummaryForTestWithWriter(Writer writer, String... args) throws Exception {
        TraceConfig traceConfig = TraceConfig.parse("java", args);
        return execute(traceConfig, writer, false);
    }

    static void run(TraceConfig traceConfig, Writer writer) throws Exception {
        execute(traceConfig, writer, true);
    }

    static TraceRunSummary executeOnce(TraceConfig traceConfig, Writer writer) throws Exception {
        return execute(traceConfig, writer, false);
    }

    private static Writer openTraceWriter(TraceConfig traceConfig) throws IOException {
        if (traceConfig.writesToStdout()) {
            return new PrintWriter(System.out, true);
        }
        Path outPath = traceConfig.getOutPath();
        Path parentPath = outPath.getParent();
        if (parentPath != null) {
            Files.createDirectories(parentPath);
        }
        return Files.newBufferedWriter(outPath, StandardCharsets.UTF_8);
    }

    private static Writer openBenchmarkWriter(TraceConfig traceConfig) throws IOException {
        if (traceConfig.writesBenchmarkToStdout()) {
            return new PrintWriter(System.out, true);
        }
        Path outPath = traceConfig.getBenchmarkOutPath();
        Path parentPath = outPath.getParent();
        if (parentPath != null) {
            Files.createDirectories(parentPath);
        }
        return Files.newBufferedWriter(outPath, StandardCharsets.UTF_8);
    }

    private static Writer openBenchmarkTraceWriter(TraceConfig traceConfig) throws IOException {
        if (!traceConfig.isTraceEnabled() || traceConfig.writesToStdout()) {
            return new DiscardingWriter();
        }
        return openTraceWriter(traceConfig);
    }

    private static void closeWriter(boolean writesToStdout, Writer writer) throws IOException {
        if (writer == null) {
            return;
        }
        if (writesToStdout) {
            writer.flush();
            return;
        }
        writer.close();
    }

    private static Region buildSingleNeuronRegion() {
        Region region = new Region("dbg");
        region.addLayer(1, 0, 0);
        return region;
    }

    private static Region buildTwoLayerRegion(double value) {
        Region region = new Region("dbg");
        int sourceLayerIndex = region.addLayer(1, 0, 0);
        int destLayerIndex = region.addLayer(1, 0, 0);
        region.connectLayers(sourceLayerIndex, destLayerIndex, 1.0, false);
        Neuron sourceNeuron = region.getLayers().get(sourceLayerIndex).getNeurons().get(0);
        primeTwoLayerSource(sourceNeuron, value);
        if (!sourceNeuron.getOutgoing().isEmpty()) {
            Weight edgeWeight = sourceNeuron.getOutgoing().get(0).getWeight();
            primeWeight(
                    edgeWeight,
                    Math.max(0.9, Math.abs(value) + 0.2),
                    Math.max(0.0, Math.abs(value) * 0.25),
                    100
            );
        }
        return region;
    }

    private static Region buildImageRegion() {
        Region region = new Region("dbg");
        Layer spatialLayer = region.getLayers().get(region.addLayer(1, 0, 0));
        Neuron spatialNeuron = spatialLayer.getNeurons().get(0);
        SlotConfig spatialConfig = spatialNeuron.getSlotConfig();
        spatialConfig.setSpatialEnabled(true);
        spatialConfig.setRowBinWidthPct(25.0);
        spatialConfig.setColBinWidthPct(25.0);
        spatialConfig.setSlotLimit(16);
        return region;
    }

    private static Region buildGrowthRegion(double value) {
        Region region = new Region("dbg");
        int layerIndex = region.addLayer(1, 0, 0);
        Layer layer = region.getLayers().get(layerIndex);
        Neuron neuron = layer.getNeurons().get(0);
        SlotConfig slotConfig = neuron.getSlotConfig();
        slotConfig.setSlotLimit(1);
        slotConfig.setFallbackGrowthThreshold(1);
        slotConfig.setNeuronGrowthCooldownTicks(0);
        slotConfig.setGrowthEnabled(true);
        slotConfig.setNeuronGrowthEnabled(true);
        slotConfig.setLayerGrowthEnabled(false);
        GrowthPolicy growthPolicy = new GrowthPolicy()
                .setEnableLayerGrowth(true)
                .setLayerCooldownTicks(0)
                .setAvgSlotsThreshold(0.0)
                .setPercentAtCapFallbackThreshold(0.0)
                .setWireProbability(1.0)
                .setNewLayerExcitatoryCount(2);
        region.setGrowthPolicy(growthPolicy);
        double baselineValue = Math.abs(value) < 1e-6 ? 0.1 : value / 2.0;
        neuron.onInput(baselineValue);
        finishTick(region, false);
        return region;
    }

    private static RegionMetrics runSingleNeuronTick(Region region, double value) {
        Layer layer = region.getLayers().get(0);
        Trace.beginTick(region.getName(), "x", value, true);
        Trace.emitLayerForward(layer, value);
        layer.forward(value);
        return finishTick(region, false);
    }

    private static RegionMetrics runTwoLayerTick(Region region, double value) {
        Layer layer = region.getLayers().get(0);
        Trace.beginTick(region.getName(), "x", value, true);
        Trace.emitLayerForward(layer, value);
        layer.forward(value);
        return finishTick(region, false);
    }

    private static RegionMetrics runImageTick(Region region, double value) {
        Layer spatialLayer = region.getLayers().get(0);
        int tickNumber = Trace.beginTick(region.getName(), "img", null, false);
        double[][] frame = new double[][]{
                {value, 0.0},
                {0.0, value / 2.0}
        };
        for (int rowIndex = 0; rowIndex < frame.length; rowIndex += 1) {
            for (int colIndex = 0; colIndex < frame[rowIndex].length; colIndex += 1) {
                double pixelValue = frame[rowIndex][colIndex];
                if (pixelValue == 0.0) {
                    continue;
                }
                final int pixelRowIndex = rowIndex;
                final int pixelColIndex = colIndex;
                final double tracePixelValue = pixelValue;
                String pixelSignalId = Trace.formatPixelSignalId(tickNumber, pixelRowIndex, pixelColIndex);
                Trace.runWithSignal(pixelSignalId, () -> {
                    Map<String, Object> eventFields = new LinkedHashMap<>();
                    Trace.addLayerFields(eventFields, spatialLayer);
                    eventFields.put("neuron_count", spatialLayer.getNeurons().size());
                    eventFields.put("input_value", tracePixelValue);
                    eventFields.put("reason", "pixel");
                    Trace.emit("layer.forward", eventFields);
                    Neuron neuron = spatialLayer.getNeurons().get(0);
                    boolean fired = neuron.onInput2D(tracePixelValue, pixelRowIndex, pixelColIndex);
                    if (fired) {
                        neuron.onOutput(tracePixelValue);
                    }
                });
            }
        }
        return finishTick(region, false);
    }

    private static RegionMetrics runGrowthTick(Region region, double value) {
        Layer layer = region.getLayers().get(0);
        Trace.beginTick(region.getName(), "x", value, true);
        Trace.emitLayerForward(layer, value);
        layer.forward(value);
        return finishTick(region, true);
    }

    private static RegionMetrics finishTick(Region region, boolean allowRegionGrowth) {
        for (Layer layer : region.getLayers()) {
            layer.endTick();
            Trace.emitBusDecay(layer);
        }
        if (allowRegionGrowth && region.getGrowthPolicy() != null) {
            try {
                GrowthEngine.maybeGrow(region, region.getGrowthPolicy());
            } catch (Throwable ignored) {
            }
        }
        RegionMetrics metrics = collectMetrics(region);
        Trace.endTick(metrics);
        return metrics;
    }

    private static RegionMetrics collectMetrics(Region region) {
        RegionMetrics metrics = new RegionMetrics();
        metrics.incDeliveredEvents();
        for (Layer layer : region.getLayers()) {
            for (Neuron neuron : layer.getNeurons()) {
                metrics.addSlots(neuron.getSlots().size());
                metrics.addSynapses(neuron.getOutgoing().size());
            }
        }
        return metrics;
    }

    private static void primeTwoLayerSource(Neuron neuron, double value) {
        neuron.onInput(value);
        Weight slot = neuron.getSlots().get(0);
        if (slot == null) {
            return;
        }
        primeWeight(
                slot,
                Math.max(0.9, Math.abs(value) + 0.2),
                Math.max(0.0, Math.abs(value) * 0.25),
                100
        );
    }

    private static TraceRunSummary execute(TraceConfig traceConfig, Writer writer, boolean emitPlainSummary) throws Exception {
        Trace.configure(traceConfig, writer);
        try {
            Region region = buildRegionForScenario(traceConfig);
            RegionMetrics metrics = new RegionMetrics();
            for (int tickOrdinal = 0; tickOrdinal < traceConfig.getTicks(); tickOrdinal += 1) {
                metrics = runScenarioTick(region, traceConfig);
                if (emitPlainSummary && !traceConfig.isTraceEnabled()) {
                    PrintWriter printWriter = new PrintWriter(writer, true);
                    printWriter.printf(
                            Locale.ROOT,
                            "scenario=%s tick=%d delivered_events=%d total_slots=%d total_synapses=%d%n",
                            traceConfig.getScenario(),
                            tickOrdinal + 1,
                            metrics.getDeliveredEvents(),
                            metrics.getTotalSlots(),
                            metrics.getTotalSynapses()
                    );
                    printWriter.flush();
                }
            }
            return summarize(region, metrics);
        } finally {
            Trace.reset();
        }
    }

    private static Region buildRegionForScenario(TraceConfig traceConfig) {
        return switch (traceConfig.getScenario()) {
            case "single_neuron" -> buildSingleNeuronRegion();
            case "two_layer" -> buildTwoLayerRegion(traceConfig.getValue());
            case "image_2d" -> buildImageRegion();
            case "growth" -> buildGrowthRegion(traceConfig.getValue());
            default -> throw new IllegalArgumentException("Unsupported scenario: " + traceConfig.getScenario());
        };
    }

    private static RegionMetrics runScenarioTick(Region region, TraceConfig traceConfig) {
        return switch (traceConfig.getScenario()) {
            case "single_neuron" -> runSingleNeuronTick(region, traceConfig.getValue());
            case "two_layer" -> runTwoLayerTick(region, traceConfig.getValue());
            case "image_2d" -> runImageTick(region, traceConfig.getValue());
            case "growth" -> runGrowthTick(region, traceConfig.getValue());
            default -> throw new IllegalArgumentException("Unsupported scenario: " + traceConfig.getScenario());
        };
    }

    private static TraceRunSummary summarize(Region region, RegionMetrics metrics) {
        int neuronCount = 0;
        for (Layer layer : region.getLayers()) {
            neuronCount += layer.getNeurons().size();
        }
        return new TraceRunSummary(
                metrics.getDeliveredEvents(),
                metrics.getTotalSlots(),
                metrics.getTotalSynapses(),
                region.getLayers().size(),
                neuronCount
        );
    }

    private static void primeWeight(
            Weight weight,
            double strengthValue,
            double thresholdValue,
            int reinforcementCount) {
        if (weight == null) {
            return;
        }
        setDoubleField(weight, "strengthValue", clamp(strengthValue, -1.0, 1.0));
        setDoubleField(weight, "thresholdValue", thresholdValue);
        setIntField(weight, "reinforcementCount", Math.max(0, reinforcementCount));
        setBooleanField(weight, "firstSeen", true);
    }

    private static double clamp(double value, double lowBound, double highBound) {
        return Math.max(lowBound, Math.min(highBound, value));
    }

    private static void setDoubleField(Object targetObject, String fieldName, double value) {
        try {
            Field field = Weight.class.getDeclaredField(fieldName);
            field.setAccessible(true);
            field.setDouble(targetObject, value);
        } catch (Throwable ignored) {
        }
    }

    private static void setIntField(Object targetObject, String fieldName, int value) {
        try {
            Field field = Weight.class.getDeclaredField(fieldName);
            field.setAccessible(true);
            field.setInt(targetObject, value);
        } catch (Throwable ignored) {
        }
    }

    private static void setBooleanField(Object targetObject, String fieldName, boolean value) {
        try {
            Field field = Weight.class.getDeclaredField(fieldName);
            field.setAccessible(true);
            field.setBoolean(targetObject, value);
        } catch (Throwable ignored) {
        }
    }

    private static final class DiscardingWriter extends Writer {
        @Override
        public void write(char[] buffer, int offset, int length) {
        }

        @Override
        public void flush() {
        }

        @Override
        public void close() {
        }
    }
}
