const std = @import("std");
const Allocator = std.mem.Allocator;
const Sigmoid = @import("activations.zig").Sigmoid;

pub const DenseLayer = struct {
    weights: []f32,
    biases: []f32,
    output: []f32,

    // Buffer per calcoli intermedi (per non riallocare memoria ogni volta)
    deltas: []f32, // L'errore di questo strato
    input_gradient: []f32, // L'errore da passare allo strato precedente

    inputs_count: usize,
    neurons_count: usize,
    allocator: Allocator,

    pub fn init(allocator: Allocator, in_size: usize, out_size: usize, seed: u64) !DenseLayer {
        const weights = try allocator.alloc(f32, in_size * out_size);
        const biases = try allocator.alloc(f32, out_size);
        const output = try allocator.alloc(f32, out_size);
        const deltas = try allocator.alloc(f32, out_size);
        const input_gradient = try allocator.alloc(f32, in_size); // Importante!

        var prng = std.Random.DefaultPrng.init(seed);
        const rand = prng.random();

        for (weights) |*w| w.* = rand.float(f32) * 2.0 - 1.0;
        for (biases) |*b| b.* = 0.0;

        return DenseLayer{
            .weights = weights,
            .biases = biases,
            .output = output,
            .deltas = deltas,
            .input_gradient = input_gradient,
            .inputs_count = in_size,
            .neurons_count = out_size,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *DenseLayer) void {
        self.allocator.free(self.weights);
        self.allocator.free(self.biases);
        self.allocator.free(self.output);
        self.allocator.free(self.deltas);
        self.allocator.free(self.input_gradient);
    }

    pub fn forward(self: *DenseLayer, input: []const f32) []f32 {
        for (0..self.neurons_count) |n| {
            var sum: f32 = 0.0;
            // Prodotto scalare (Input * Pesi)
            for (0..self.inputs_count) |i| {
                sum += input[i] * self.weights[n * self.inputs_count + i];
            }
            // Aggiungiamo bias e applichiamo Sigmoide
            self.output[n] = Sigmoid.apply(sum + self.biases[n]);
        }
        return self.output;
    }

    // Questa è la parte magica
    pub fn backward(self: *DenseLayer, next_layer_gradients: []const f32, prev_layer_input: []const f32, lr: f32) []f32 {
        // 1. Calcoliamo i "Delta" (Errore locale * derivata attivazione)
        for (0..self.neurons_count) |n| {
            const derivative = Sigmoid.derivative(self.output[n]);
            self.deltas[n] = next_layer_gradients[n] * derivative;
        }

        // 2. Calcoliamo il gradiente da passare indietro (Input Gradient)
        // Questo serve allo strato PRECEDENTE per correggersi
        @memset(self.input_gradient, 0.0);
        for (0..self.inputs_count) |i| {
            for (0..self.neurons_count) |n| {
                // Sommiamo il contributo di ogni neurone connesso a questo input
                self.input_gradient[i] += self.deltas[n] * self.weights[n * self.inputs_count + i];
            }
        }

        // 3. Aggiorniamo i pesi e i bias (Gradient Descent)
        for (0..self.neurons_count) |n| {
            for (0..self.inputs_count) |i| {
                // Peso -= LearningRate * Delta * InputOriginale
                self.weights[n * self.inputs_count + i] -= lr * self.deltas[n] * prev_layer_input[i];
            }
            self.biases[n] -= lr * self.deltas[n];
        }

        return self.input_gradient;
    }
};
