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, use_sigmoid: bool) !void {
        const layer = try DenseLayer.init(self.allocator, input_size, output_size, seed, use_sigmoid);
        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 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;

            const new_gradients = layer.backward(next_gradients, prev_input, lr);

            if (next_gradients.ptr != output_errors.ptr) {
                self.allocator.free(next_gradients);
            }

            next_gradients = new_gradients;
        }

        if (next_gradients.ptr != output_errors.ptr) {
            self.allocator.free(next_gradients);
        }

        return total_loss;
    }

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

        const MagicNumber: u64 = 0xDEADBEEF;

        try file.writeAll(std.mem.asBytes(&MagicNumber));

        const layer_count = @as(u64, self.layers.items.len);
        try file.writeAll(std.mem.asBytes(&layer_count));

        for (self.layers.items) |layer| {
            const inputs = @as(u64, layer.inputs_count);
            const neurons = @as(u64, layer.neurons_count);
            try file.writeAll(std.mem.asBytes(&inputs));
            try file.writeAll(std.mem.asBytes(&neurons));

            const weights_bytes = std.mem.sliceAsBytes(layer.weights);
            try file.writeAll(weights_bytes);

            const biases_bytes = std.mem.sliceAsBytes(layer.biases);
            try file.writeAll(biases_bytes);
        }
    }

    pub fn load(allocator: Allocator, file_path: []const u8) !Network {
        const file = try std.fs.cwd().openFile(file_path, .{});
        defer file.close();

        var net = Network.init(allocator);
        var success = false;

        defer if (!success) net.deinit();

        var magic: u64 = 0;
        _ = try file.readAll(std.mem.asBytes(&magic));
        if (magic != 0xDEADBEEF) return error.InvalidNetworkFile;

        var layer_count: u64 = 0;
        _ = try file.readAll(std.mem.asBytes(&layer_count));

        var i: u64 = 0;
        while (i < layer_count) : (i += 1) {
            var inputs: u64 = 0;
            var neurons: u64 = 0;

            _ = try file.readAll(std.mem.asBytes(&inputs));
            _ = try file.readAll(std.mem.asBytes(&neurons));

            const is_last = (i == layer_count - 1);
            try net.addLayer(@intCast(inputs), @intCast(neurons), 0, !is_last);

            const layer = &net.layers.items[net.layers.items.len - 1];

            const weights_bytes = std.mem.sliceAsBytes(layer.weights);
            _ = try file.readAll(weights_bytes);

            const biases_bytes = std.mem.sliceAsBytes(layer.biases);
            _ = try file.readAll(biases_bytes);
        }

        success = true;
        return net;
    }

    pub fn exportJSON(self: *Network, file_path: []const u8, epoch: usize, loss: f32, input_snapshot: ?[]const 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", .{ epoch, loss });

        if (input_snapshot) |pixels| {
            try file.writeAll("  \"input_pixels\": [");
            for (pixels, 0..) |p, idx| {
                try Utils.print(file, "{d:.2}", .{p});
                if (idx < pixels.len - 1) try file.writeAll(",");
            }
            try file.writeAll("],\n");
        }

        try file.writeAll("  \"layers\": [\n");

        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 });

            const max_w = if (layer.weights.len > 1000) 100 else layer.weights.len;

            for (layer.weights[0..max_w], 0..) |w, w_idx| {
                try Utils.print(file, "{d:.4}", .{w});
                if (w_idx < max_w - 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");
    }
};
