<?php

namespace App\Models\NetworksTraining;

use App\Models\ActivationsFunctions;
use App\Models\Perceptrons\GradientDescentPerceptron;
use App\Models\Perceptrons\NetworkPerceptron;
use App\Models\Perceptrons\Perceptron;
use App\Models\Perceptrons\SigmoidPerceptron;
use App\Services\DatasetReader\IDataSetReader;
use App\Services\IterationEventBuffer\IPerceptronIterationEventBuffer;
use App\Services\SynapticWeightsProvider\ISynapticWeightsProvider;
use App\Services\SynapticWeightsProvider\SimpleNetworkWeightsProvider;
use Illuminate\Support\Arr;

class MultiLayerPerceptronTraining extends NetworkTraining
{
    private Perceptron $network;

    private array $labels;

    private bool $isRegression;

    public ActivationsFunctions $activationFunction = ActivationsFunctions::SIGMOID;
    public ?ActivationsFunctions $presentationLayerActivationFunction = ActivationsFunctions::STEP;

    private float $epochError;

    public function __construct(
        IDataSetReader $datasetReader,
        protected float $learningRate,
        int $maxEpochs,
        protected int $hiddenLayers,
        protected int $hiddenLayersNeurons,
        ISynapticWeightsProvider $synapticWeightsProvider,
        IPerceptronIterationEventBuffer $iterationEventBuffer,
        string $sessionId,
        string $trainingId,
        private float $minError,
    ) {
        parent::__construct($datasetReader, $maxEpochs, $iterationEventBuffer, $sessionId, $trainingId);
        $this->labels = $datasetReader->getLabels();
        $this->isRegression = $datasetReader->getOutputSize() === 1
            || ($datasetReader->getOutputSize() > 2
                && $datasetReader->getOutputSize() * 2 >= $datasetReader->getEpochExamplesCount());
        if ($this->isRegression) {
            $this->activationFunction = ActivationsFunctions::LINEAR;
        }

        $networkWeightsProvider = new SimpleNetworkWeightsProvider($synapticWeightsProvider);
        $this->network = new NetworkPerceptron(
            $networkWeightsProvider->generate(
                $datasetReader->getInputSize(),
                $this->isRegression ? 1 : $datasetReader->getOutputSize(),
                $this->hiddenLayers,
                $this->hiddenLayersNeurons,
            ),
            $datasetReader->getInputSize(),
            SigmoidPerceptron::class,
            GradientDescentPerceptron::class,
        );
    }

    public function start(): void
    {
        $this->epoch = 0;
        do {
            $this->epochError = 0;
            $this->epoch++;

            $inputsForCurrentEpoch = [];

            while ($nextRow = $this->datasetReader->getNextLine()) {
                $inputsForCurrentEpoch[] = $nextRow;
                $inputs = array_slice($nextRow, 0, -1);
                $correctOutput = (float) end($nextRow);

                $iterationError = $this->iterationFunction($inputs, $correctOutput);

                // Synaptic weights correction after each example
                $synaptic_weights = $this->network->getSynapticWeights();
                $inputs_with_bias = array_merge([1], $inputs); // Add bias input

                // Updates the weights
                $this->network->setSynapticWeights(
                    $this->getUpdatedSynapticWeights($synaptic_weights, $iterationError, $inputs_with_bias)
                );

                // Broadcast the training iteration event
                $this->addIterationToBuffer(array_sum($iterationError), $this->network->getSynapticWeights());

                // $this->iterationEventBuffer->flush();
            }

            // Calculte the average error for the epoch with the last synaptic weights
            foreach ($inputsForCurrentEpoch as $inputsWithLabel) {
                $inputs = array_slice($inputsWithLabel, 0, -1);
                $correctOutput = (float) end($inputsWithLabel);
                $iterationError = $this->iterationFunction($inputs, $correctOutput);
                foreach ($iterationError as $error) {
                    $this->epochError += ($error ** 2) / 2; // Squared error for the example
                }
            }
            $this->epochError /= $this->datasetReader->getEpochExamplesCount(); // Average error for the epoch

            $this->datasetReader->reset(); // Reset the dataset for the next iteration
        } while ($this->epoch < $this->maxEpochs && ! $this->stopCondition());

        $this->iterationEventBuffer->flush(); // Ensure all iterations are sent to the frontend

        $this->checkPassedMaxIterations($this->epochError);
    }

    protected function stopCondition(): bool
    {
        $condition = $this->epochError <= $this->minError;
        if ($condition === true) {
            $this->broadcastTrainingEnded('Le perceptron à atteint l\'erreur minimale');
        }

        return $condition;
    }

    private function iterationFunction(array $inputs, float $correctOutput): array
    {
        $outputs = $this->network->test($inputs);
        $desiredOutput = $this->getDesiredOutputFromCorrectOutput($correctOutput);

        $errors = [];
        foreach ($outputs as $index => $output) {
            $error = $desiredOutput[$index] - $output;
            $errors[] = $error;
        }

        return $errors;
    }

    /**
     * Backpropagation  of error gradients to update synaptic weights.
     *
     */
    private function getUpdatedSynapticWeights(array $synaptic_weights, array $iterationError, array $inputs): array
    {
        $layerInputs = [$inputs];

        // Reproduce NetworkPerceptron::test() using each neuron's activation function.
        foreach ($synaptic_weights as $layerIndex => $layerWeights) {
            $previousLayerOutputs = $layerInputs[array_key_last($layerInputs)];
            $layerOutputs = [];

            foreach ($layerWeights as $neuronIndex => $neuronWeights) {
                $weightedSum = array_sum(array_map(
                    fn ($input, $weight): float => $input * $weight,
                    $previousLayerOutputs,
                    $neuronWeights,
                ));

                $neuron = $this->network->network[$layerIndex + 1][$neuronIndex];
                $layerOutputs[] = $neuron->activationFunction($weightedSum);
            }

            $layerInputs[] = array_merge([1], $layerOutputs);
        }

        $deltas = array_fill(0, count($synaptic_weights), []);
        $lastLayerIndex = count($synaptic_weights) - 1;

        // Output delta includes the output neuron's activation derivative.
        foreach ($iterationError as $neuronIndex => $error) {
            $neuron = $this->network->network[$lastLayerIndex + 1][$neuronIndex];
            $output = $layerInputs[$lastLayerIndex + 1][$neuronIndex + 1];

            $deltas[$lastLayerIndex][$neuronIndex] =
                $error * $this->activationDerivative($neuron, $output);
        }

        // Hidden-layer deltas use the original, unchanged weights.
        for ($layerIndex = $lastLayerIndex - 1; $layerIndex >= 0; $layerIndex--) {
            foreach ($synaptic_weights[$layerIndex] as $neuronIndex => $unusedNeuronWeights) {
                $nextLayerDelta = 0.0;

                foreach ($synaptic_weights[$layerIndex + 1] as $nextNeuronIndex => $nextNeuronWeights) {
                    // Index zero is the next layer's bias weight.
                    $nextLayerDelta +=
                        $nextNeuronWeights[$neuronIndex + 1]
                        * $deltas[$layerIndex + 1][$nextNeuronIndex];
                }

                $neuron = $this->network->network[$layerIndex + 1][$neuronIndex];
                $output = $layerInputs[$layerIndex + 1][$neuronIndex + 1];

                $deltas[$layerIndex][$neuronIndex] =
                    $nextLayerDelta * $this->activationDerivative($neuron, $output);
            }
        }

        $updatedWeights = [];

        foreach ($synaptic_weights as $layerIndex => $layerWeights) {
            $updatedLayerWeights = [];

            foreach ($layerWeights as $neuronIndex => $neuronWeights) {
                $updatedLayerWeights[] = array_map(
                    fn ($weight, $weightIndex): float => $weight
                        + $this->learningRate
                        * $deltas[$layerIndex][$neuronIndex]
                        * $layerInputs[$layerIndex][$weightIndex],
                    $neuronWeights,
                    array_keys($neuronWeights),
                );
            }

            $updatedWeights[] = $updatedLayerWeights;
        }

        return $updatedWeights;
    }

    private function activationDerivative(Perceptron $neuron, float $output): float
    {
        // Numerical derivative works for any activationFunction implementation.
        // $epsilon = 1e-6;

        // The activation function requires weighted input, which is unavailable here.
        // For sigmoid, use its derivative directly.
        if ($neuron instanceof SigmoidPerceptron) {
            return $output * (1 - $output);
        }

        // Linear activation, used commonly by gradient-descent output neurons.
        return 1.0;
    }

    private function getDesiredOutputFromCorrectOutput(float $correctOutput): array
    {
        if ($this->isRegression) {
            return [$correctOutput];
        }

        $desiredOutput = array_fill(0, count($this->labels), 0);
        $labelIndex = Arr::first(
            array_keys($this->labels),
            fn ($key) => $this->labels[$key] == $correctOutput
        );

        if ($labelIndex !== null) {
            $desiredOutput[$labelIndex] = 1;
        }

        return $desiredOutput;
    }

    public function getSynapticWeights(): array
    {
        return [[$this->network->getSynapticWeights()]];
    }
}