module systolic_array_NxN #(
    parameter WIDTH = 32,
    parameter ROWS = 16,
    parameter COLS = 16
)(
    input wire clk,
    input wire rst_n,
    input wire [WIDTH-1:0] a_in [ROWS-1:0],
    input wire [WIDTH-1:0] b_in [COLS-1:0],
    output wire [WIDTH-1:0] c_out [ROWS-1:0][COLS-1:0]
);

    // Wire matrix
    wire [WIDTH-1:0] h_wires [ROWS:0][COLS:0]; // Horizontal
    wire [WIDTH-1:0] v_wires [ROWS:0][COLS:0]; // Vertical

    // Assign inputs
    genvar r, c;
    generate
        for (r = 0; r < ROWS; r = r + 1) begin
            assign h_wires[r][0] = a_in[r];
        end
        for (c = 0; c < COLS; c = c + 1) begin
            assign v_wires[0][c] = b_in[c];
        end

        for (r = 0; r < ROWS; r = r + 1) begin : rows
            for (c = 0; c < COLS; c = c + 1) begin : cols
                wire [WIDTH-1:0] c_acc_wire;

                systolic_pe pe (
                    .clk(clk),
                    .rst_n(rst_n),
                    .a_in(h_wires[r][c]),
                    .b_in(v_wires[r][c]),
                    .a_out(h_wires[r][c+1]),
                    .b_out(v_wires[r+1][c]),
                    .c_acc(c_acc_wire)
                );

                assign c_out[r][c] = c_acc_wire;
            end
        end
    endgenerate

endmodule
