const std = @import("std");

pub const SimpleNetwork = struct {
    w_hidden: [2][2]f32,
    b_hidden: [2]f32,

    w_output: [2]f32,
    b_output: f32,

    hidden_outputs: [2]f32,

    pub fn init(seed: u64) SimpleNetwork {
        var prng = std.Random.DefaultPrng.init(seed);
        const rand = prng.random();

        var net = SimpleNetwork{
            .w_hidden = undefined,
            .b_hidden = undefined,
            .w_output = undefined,
            .b_output = 0.0,
            .hidden_outputs = undefined,
        };

        for (&net.w_hidden) |*row| {
            row[0] = rand.float(f32) * 2.0 - 1.0;
            row[1] = rand.float(f32) * 2.0 - 1.0;
        }
        for (&net.b_hidden) |*b| b.* = 0.0;

        for (&net.w_output) |*w| w.* = rand.float(f32) * 2.0 - 1.0;

        return net;
    }

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

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

    pub fn forward(self: *SimpleNetwork, input: [2]f32) f32 {
        for (0..2) |i| {
            const sum = (input[0] * self.w_hidden[i][0]) +
                (input[1] * self.w_hidden[i][1]) +
                self.b_hidden[i];
            self.hidden_outputs[i] = sigmoid(sum);
        }

        const sum_out = (self.hidden_outputs[0] * self.w_output[0]) +
            (self.hidden_outputs[1] * self.w_output[1]) +
            self.b_output;

        return sigmoid(sum_out);
    }

    pub fn train(self: *SimpleNetwork, input: [2]f32, target: f32, lr: f32) f32 {
        const prediction = self.forward(input);

        const output_error = target - prediction;
        const output_delta = output_error * sigmoid_derivative(prediction);

        var hidden_deltas: [2]f32 = undefined;
        for (0..2) |i| {
            const error_contrib = output_delta * self.w_output[i];
            hidden_deltas[i] = error_contrib * sigmoid_derivative(self.hidden_outputs[i]);
        }

        for (0..2) |i| {
            self.w_output[i] += lr * output_delta * self.hidden_outputs[i];
        }
        self.b_output += lr * output_delta;

        for (0..2) |i| {
            self.w_hidden[i][0] += lr * hidden_deltas[i] * input[0];
            self.w_hidden[i][1] += lr * hidden_deltas[i] * input[1];
            self.b_hidden[i] += lr * hidden_deltas[i];
        }

        return output_error * output_error;
    }
};
