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

const SimdWidth = 8;
const Vec = @Vector(SimdWidth, f32);

pub const DenseLayer = struct {
    weights: []f32,
    biases: []f32,
    output: []f32,
    inputs_count: usize,
    neurons_count: usize,
    use_sigmoid: bool,
    allocator: Allocator,

    fn sigmoid(x: f32) f32 {
        return 1.0 / (1.0 + std.math.exp(-x));
    }

    fn sigmoidDerivative(x: f32) f32 {
        return x * (1.0 - x);
    }

    pub fn init(allocator: Allocator, inputs: usize, neurons: usize, seed: u64, use_sigmoid: bool) !DenseLayer {
        const weights = try allocator.alloc(f32, inputs * neurons);
        const biases = try allocator.alloc(f32, neurons);
        const output = try allocator.alloc(f32, neurons);

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

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

        return DenseLayer{
            .weights = weights,
            .biases = biases,
            .output = output,
            .inputs_count = inputs,
            .neurons_count = neurons,
            .use_sigmoid = use_sigmoid,
            .allocator = allocator,
        };
    }

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

    pub fn forward(self: *DenseLayer, input: []const f32) []const f32 {
        for (0..self.neurons_count) |n| {
            var sum: f32 = self.biases[n];
            const w_start = n * self.inputs_count;

            var vec_sum: Vec = @splat(0.0);
            var i: usize = 0;
            while (i + SimdWidth <= self.inputs_count) : (i += SimdWidth) {
                const v_in: Vec = input[i..][0..SimdWidth].*;
                const v_w: Vec = self.weights[w_start + i ..][0..SimdWidth].*;
                vec_sum += v_in * v_w;
            }
            sum += @reduce(.Add, vec_sum);

            while (i < self.inputs_count) : (i += 1) {
                sum += input[i] * self.weights[w_start + i];
            }

            if (self.use_sigmoid) {
                self.output[n] = sigmoid(sum);
            } else {
                self.output[n] = sum;
            }
        }
        return self.output;
    }

    pub fn backward(self: *DenseLayer, output_gradient: []const f32, input_vals: []const f32, learning_rate: f32) []f32 {
        const input_gradient = self.allocator.alloc(f32, self.inputs_count) catch @panic("OOM");
        @memset(input_gradient, 0.0);

        for (0..self.neurons_count) |n| {
            var derivative: f32 = 1.0;
            if (self.use_sigmoid) {
                derivative = sigmoidDerivative(self.output[n]);
            }

            const delta = output_gradient[n] * derivative;
            const w_start = n * self.inputs_count;

            self.biases[n] -= learning_rate * delta;

            const v_delta: Vec = @splat(delta);
            const v_lr: Vec = @splat(learning_rate);
            const v_change_factor = v_delta * v_lr;

            var i: usize = 0;
            while (i + SimdWidth <= self.inputs_count) : (i += SimdWidth) {
                var v_w: Vec = self.weights[w_start + i ..][0..SimdWidth].*;
                const v_in: Vec = input_vals[i..][0..SimdWidth].*;

                var v_in_grad: Vec = input_gradient[i..][0..SimdWidth].*;
                v_in_grad += v_w * v_delta;
                input_gradient[i..][0..SimdWidth].* = v_in_grad;

                v_w -= v_in * v_change_factor;
                self.weights[w_start + i ..][0..SimdWidth].* = v_w;
            }

            while (i < self.inputs_count) : (i += 1) {
                const w_idx = w_start + i;
                const old_weight = self.weights[w_idx];
                input_gradient[i] += old_weight * delta;
                self.weights[w_idx] -= input_vals[i] * delta * learning_rate;
            }
        }
        return input_gradient;
    }
};
