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

pub const Network = struct {
    layers: std.ArrayList(DenseLayer),
    allocator: Allocator,

    pub fn init(allocator: Allocator) Network {
        return Network{
            .layers = std.ArrayList(DenseLayer){},
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Network) void {
        for (self.layers.items) |*layer| {
            layer.deinit();
        }
        self.layers.deinit(self.allocator);
    }

    pub fn addLayer(self: *Network, input_size: usize, output_size: usize, seed: u64) !void {
        const layer = try DenseLayer.init(self.allocator, input_size, output_size, seed);
        try self.layers.append(self.allocator, layer);
    }

    pub fn forward(self: *Network, input: []const f32) []const f32 {
        var current_input = input;
        for (self.layers.items) |*layer| {
            current_input = layer.forward(current_input);
        }
        return current_input;
    }

    pub fn printTopology(self: *Network) void {
        std.debug.print("Architettura Rete: [Input]", .{});
        for (self.layers.items) |layer| {
            std.debug.print(" -> [Dense:{d}]", .{layer.neurons_count});
        }
        std.debug.print("\n", .{});
    }

    pub fn exportJSON(self: *Network, file_path: []const u8, epoch: usize, loss: f32) !void {
        const file = try std.fs.cwd().createFile(file_path, .{});
        defer file.close();

        const Utils = struct {
            fn print(f: std.fs.File, comptime fmt: []const u8, args: anytype) !void {
                var buf: [256]u8 = undefined;
                const text = try std.fmt.bufPrint(&buf, fmt, args);
                try f.writeAll(text);
            }
        };

        try Utils.print(file, "{{\n  \"epoch\": {d},\n  \"loss\": {d:.6},\n  \"layers\": [\n", .{ epoch, loss });

        for (self.layers.items, 0..) |layer, i| {
            try Utils.print(file, "    {{\n      \"layer_index\": {d},\n      \"neurons\": {d},\n      \"inputs\": {d},\n      \"weights\": [", .{ i, layer.neurons_count, layer.inputs_count });

            for (layer.weights, 0..) |w, w_idx| {
                try Utils.print(file, "{d:.4}", .{w});
                if (w_idx < layer.weights.len - 1) try file.writeAll(", ");
            }

            try file.writeAll("]\n    }");
            if (i < self.layers.items.len - 1) try file.writeAll(",\n");
        }

        try file.writeAll("\n  ]\n}\n");
    }

    pub fn train(self: *Network, input: []const f32, target: []const f32, lr: f32) !f32 {
        _ = self.forward(input);

        const last_layer_idx = self.layers.items.len - 1;
        const last_layer = &self.layers.items[last_layer_idx];

        var output_errors = try self.allocator.alloc(f32, last_layer.neurons_count);
        defer self.allocator.free(output_errors);

        var total_loss: f32 = 0.0;
        for (0..last_layer.neurons_count) |i| {
            const err = last_layer.output[i] - target[i];
            output_errors[i] = err;
            total_loss += err * err;
        }

        var next_gradients = output_errors;
        var i: usize = self.layers.items.len;
        while (i > 0) {
            i -= 1;
            var layer = &self.layers.items[i];
            const prev_input = if (i == 0) input else self.layers.items[i - 1].output;
            next_gradients = layer.backward(next_gradients, prev_input, lr);
        }

        return total_loss;
    }
};
