<script setup lang="ts">
import type {
    ChartDataset,
    ChartTypeRegistry,
    BubbleDataPoint,
    Point,
} from 'chart.js';
import { computed, ref } from 'vue';
import { Chart } from 'vue-chartjs';
import { colors, gridColor, gridColorBold } from '@/types/graphs';
import type { Iteration } from '@/types/perceptron';

const props = defineProps<{
    cleanedDataset: { label: number; data: { x: number; y: number }[] }[];
    iterations: Iteration[];
    activationFunction: (x: number) => number;
}>();

const examplesNumber = computed(() => {
    return props.cleanedDataset.reduce((sum, dataset) => sum + dataset.data.length, 0);
});

const farLeftDataPointX = computed(() => {
    if (props.cleanedDataset.length === 0) {
        return 0;
    }
    const minX = Math.min(
        ...props.cleanedDataset.flatMap((d) => d.data.map((point) => point.x)),
    );
    return minX;
});
const farBottomDataPointY = computed(() => {
    if (props.cleanedDataset.length === 0) {
        return 0;
    }
    const minY = Math.min(
        ...props.cleanedDataset.flatMap((d) => d.data.map((point) => point.y)),
    );
    return minY;
});
const farRightDataPointX = computed(() => {
    if (props.cleanedDataset.length === 0) {
        return 0;
    }
    const maxX = Math.max(
        ...props.cleanedDataset.flatMap((d) => d.data.map((point) => point.x)),
    );
    return maxX;
});
const farTopDataPointY = computed(() => {
    if (props.cleanedDataset.length === 0) {
        return 0;
    }
    const maxY = Math.max(
        ...props.cleanedDataset.flatMap((d) => d.data.map((point) => point.y)),
    );
    return maxY;
});

function getPerceptronOutput(
    weightsNetwork: number[][][],
    inputs: number[],
): number[] {
    for (const layer of weightsNetwork) {
        const nextInputs: number[] = [];

        for (const neuron of layer) {
            const bias = neuron[0];
            const weights = neuron.slice(1);

            let sum = bias;

            for (let i = 0; i < weights.length; i++) {
                sum += weights[i] * inputs[i];
            }

            const activated = props.activationFunction(sum);

            nextInputs.push(activated);
        }

        inputs = nextInputs;
    }
    return inputs;
}

const nonLinearGraph = ref<boolean>(false);
function getPerceptronDecisionBoundaryDataset(
    networkWeights: number[][][],
    activationFunction: (x: number) => number = (x) => x,
): ChartDataset<
    keyof ChartTypeRegistry,
    number | Point | [number, number] | BubbleDataPoint | null
>[] {
    const label = 'Ligne de décision du Perceptron';
    console.log('Calculating decision boundary with weights:', networkWeights);

    if (
        networkWeights.length == 1 &&
        networkWeights[0].length == 1 &&
        networkWeights[0][0].length <= 3
    ) {
        nonLinearGraph.value = false;
        // Unique, 3 weights perceptron
        const perceptronWeights = [...networkWeights[0][0]]; // Copy of the unique perceptron weights

        function perceptronLine(x: number): number {
            if (perceptronWeights.length < 3) {
                // If we have less than 3 weights, we assume missing weights are zero
                return getPerceptronOutput(networkWeights, [x])[0];
            }

            // w0 + w1*x + w2*y = 0 => y = -(w1/w2)*x - w0/w2
            const w2 = perceptronWeights[2] == 0 ? 1e-6 : perceptronWeights[2]; // Avoid division by zero
            return -(perceptronWeights[1] / w2) * x - perceptronWeights[0] / w2;
        }

        // Simple line
        return [
            {
                type: 'line',
                label: label,
                data: [
                    {
                        x: farLeftDataPointX.value - 1,
                        y: perceptronLine(farLeftDataPointX.value - 1),
                    },
                    {
                        x: farRightDataPointX.value + 1,
                        y: perceptronLine(farRightDataPointX.value + 1),
                    },
                ],
                borderColor: '#FFF',
                borderWidth: 2,
                pointRadius: 0,
            },
        ];
    } else {
        nonLinearGraph.value = true;

        const bubbleTransparency = '30';
        const isInDataThreshold = 0.0;


        // -------- 1️⃣ Construction des datasets --------
        const datasets: {
            type: string;
            label: string;
            data: Point[];
            backgroundColor: string;
            pointRadius: number;
            borderWidth: number;
            order: number;
        }[] = [];
        // For the number of neuron in the last layer
        const lastLayer = networkWeights[networkWeights.length - 1];
        for (let i = 0; i < lastLayer.length; i++) {
            const dataset = {
                type: 'scatter',
                label: label,
                data: [], // Will be filled with the decision boundary points
                backgroundColor: colors[i] + bubbleTransparency || '#AAA',
                pointRadius: 15,
                borderWidth: 0,
                order: -1,
            };
            datasets.push(dataset);
        }

        // -------- 2️⃣ Échantillonnage grille --------
        const step =
            Math.abs(
                farRightDataPointX.value + 1 - (farLeftDataPointX.value - 1),
            ) / 50;

        for (
            let x = farLeftDataPointX.value - 1;
            x <= farRightDataPointX.value + 1;
            x += step
        ) {
            for (
                let y = farBottomDataPointY.value - 1;
                y <= farTopDataPointY.value + 1;
                y += step
            ) {
                const values = getPerceptronOutput(networkWeights, [x, y]);
                values.forEach((v, i) => {
                    if (v > isInDataThreshold) {
                        datasets[i].data.push({ x, y });
                    }
                });
            }
        }

        // -------- 3️⃣ Dataset ChartJS --------
        return datasets;
    }
}
</script>

<template>
    <Chart
        v-if="props.cleanedDataset.length > 0 || props.iterations.length > 0"
        class="flex bg-primary dark:bg-transparent!"
        type="scatter"
        :options="{
            responsive: true,
            maintainAspectRatio: true,
            plugins: {
                legend: {
                    position: 'top',
                },
                title: {
                    display: true,
                    text: 'Ligne de décision du Perceptron',
                },
            },
            animation: {
                duration: nonLinearGraph || examplesNumber > 10 ? 0 : 1000, // Disable animations for instant updates
            },
            layout: {
                padding: {
                    left: 10,
                    right: 10,
                    top: 10,
                    bottom: 10,
                },
            },
            scales: {
                x: {
                    type: 'linear',
                    position: 'bottom',
                    grid: {
                        color: function (context) {
                            if (context.tick.value == 0) {
                                return gridColorBold;
                            }

                            return gridColor;
                        },
                    },
                },
                y: {
                    type: 'linear',
                    position: 'left',
                    grid: {
                        color: function (context) {
                            if (context.tick.value == 0) {
                                return gridColorBold;
                            }

                            return gridColor;
                        },
                    },
                },
            },
        }"
        :data="{
            datasets: [
                // Points from the dataset
                ...props.cleanedDataset.map((dataset, index) => ({
                    type: 'scatter',
                    label: `Label ${dataset.label}`,
                    data: dataset.data,
                    backgroundColor: colors[index] || '#AAA',
                })),

                // Perceptron decision boundary
                ...getPerceptronDecisionBoundaryDataset(
                    props.iterations.length > 0
                        ? props.iterations[props.iterations.length - 1].weights
                        : [[[0, 0, 0]]],
                    props.activationFunction,
                ),
            ],
        }"
    />
</template>
