module systolic_pe #(
    parameter WIDTH = 32
)(
    input wire clk,
    input wire rst_n,
    input wire [WIDTH-1:0] a_in,
    input wire [WIDTH-1:0] b_in,
    output reg [WIDTH-1:0] a_out,
    output reg [WIDTH-1:0] b_out,
    output reg [WIDTH-1:0] c_acc // Accumulator
);

    // Tropical PE: C = max(C, A + B)
    // Data flow: A flows right, B flows down.

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            a_out <= 0;
            b_out <= 0;
            c_acc <= 0;
        end else begin
            // Forward Data
            a_out <= a_in;
            b_out <= b_in;

            // Tropical MAC: C = max(C, A+B)
            if ((a_in + b_in) > c_acc) begin
                c_acc <= a_in + b_in;
            end
        end
    end

endmodule

module systolic_array_4x4 #(
    parameter WIDTH = 32
)(
    input wire clk,
    input wire rst_n,
    // Inputs: 4 rows of A, 4 cols of B
    input wire [WIDTH-1:0] a0, a1, a2, a3,
    input wire [WIDTH-1:0] b0, b1, b2, b3,
    // Outputs: 16 Accumulators (Full C matrix)
    output wire [WIDTH-1:0] c00, c01, c02, c03,
    output wire [WIDTH-1:0] c10, c11, c12, c13,
    output wire [WIDTH-1:0] c20, c21, c22, c23,
    output wire [WIDTH-1:0] c30, c31, c32, c33
);

    // Wire matrix for internal connections
    wire [WIDTH-1:0] h_wires [4:0][4:0]; // Horizontal (A)
    wire [WIDTH-1:0] v_wires [4:0][4:0]; // Vertical (B)

    // Assign inputs
    assign h_wires[0][0] = a0; assign h_wires[1][0] = a1; assign h_wires[2][0] = a2; assign h_wires[3][0] = a3;
    assign v_wires[0][0] = b0; assign v_wires[0][1] = b1; assign v_wires[0][2] = b2; assign v_wires[0][3] = b3;

    // Generate Grid
    genvar i, j;
    generate
        for (i = 0; i < 4; i = i + 1) begin : rows
            for (j = 0; j < 4; j = j + 1) begin : cols
                // Map C outputs
                wire [WIDTH-1:0] c_val;

                systolic_pe pe (
                    .clk(clk),
                    .rst_n(rst_n),
                    .a_in(h_wires[i][j]),
                    .b_in(v_wires[i][j]),
                    .a_out(h_wires[i][j+1]),
                    .b_out(v_wires[i+1][j]),
                    .c_acc(c_val)
                );
            end
        end
    endgenerate

    // Manual mapping for output wires (Simpler for simulation probing)
    // Note: Verilog generate creates scopes, accessing them directly is tricky without SystemVerilog interfaces.
    // For this prototype, we rely on the internal state or would need to fan out registers.
    // Assuming 'c_acc' needs to be tapped. Ideally we shift it out.
    // For V3 perfect implementation, let's just expose the accumulators directly from the instances?
    // Verilog doesn't allow cross-module peeking easily.
    // Let's assume a readout mechanism or just use hierarchichal paths in simulation.
    // BUT, to be "perfect", let's instantiate explicitly to wire outputs.

    systolic_pe pe00 (.clk(clk), .rst_n(rst_n), .a_in(a0), .b_in(b0), .a_out(h_wires[0][1]), .b_out(v_wires[1][0]), .c_acc(c00));
    systolic_pe pe01 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[0][1]), .b_in(b1), .a_out(h_wires[0][2]), .b_out(v_wires[1][1]), .c_acc(c01));
    systolic_pe pe02 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[0][2]), .b_in(b2), .a_out(h_wires[0][3]), .b_out(v_wires[1][2]), .c_acc(c02));
    systolic_pe pe03 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[0][3]), .b_in(b3), .a_out(h_wires[0][4]), .b_out(v_wires[1][3]), .c_acc(c03));

    // Row 1 ... (Full Explicit Instantiation for robustness)
    systolic_pe pe10 (.clk(clk), .rst_n(rst_n), .a_in(a1), .b_in(v_wires[1][0]), .a_out(h_wires[1][1]), .b_out(v_wires[2][0]), .c_acc(c10));
    systolic_pe pe11 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[1][1]), .b_in(v_wires[1][1]), .a_out(h_wires[1][2]), .b_out(v_wires[2][1]), .c_acc(c11));
    systolic_pe pe12 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[1][2]), .b_in(v_wires[1][2]), .a_out(h_wires[1][3]), .b_out(v_wires[2][2]), .c_acc(c12));
    systolic_pe pe13 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[1][3]), .b_in(v_wires[1][3]), .a_out(h_wires[1][4]), .b_out(v_wires[2][3]), .c_acc(c13));

    // Row 2
    systolic_pe pe20 (.clk(clk), .rst_n(rst_n), .a_in(a2), .b_in(v_wires[2][0]), .a_out(h_wires[2][1]), .b_out(v_wires[3][0]), .c_acc(c20));
    systolic_pe pe21 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[2][1]), .b_in(v_wires[2][1]), .a_out(h_wires[2][2]), .b_out(v_wires[3][1]), .c_acc(c21));
    systolic_pe pe22 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[2][2]), .b_in(v_wires[2][2]), .a_out(h_wires[2][3]), .b_out(v_wires[3][2]), .c_acc(c22));
    systolic_pe pe23 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[2][3]), .b_in(v_wires[2][3]), .a_out(h_wires[2][4]), .b_out(v_wires[3][3]), .c_acc(c23));

    // Row 3
    systolic_pe pe30 (.clk(clk), .rst_n(rst_n), .a_in(a3), .b_in(v_wires[3][0]), .a_out(h_wires[3][1]), .b_out(v_wires[4][0]), .c_acc(c30));
    systolic_pe pe31 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[3][1]), .b_in(v_wires[3][1]), .a_out(h_wires[3][2]), .b_out(v_wires[4][1]), .c_acc(c31));
    systolic_pe pe32 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[3][2]), .b_in(v_wires[3][2]), .a_out(h_wires[3][3]), .b_out(v_wires[4][2]), .c_acc(c32));
    systolic_pe pe33 (.clk(clk), .rst_n(rst_n), .a_in(h_wires[3][3]), .b_in(v_wires[3][3]), .a_out(h_wires[3][4]), .b_out(v_wires[4][3]), .c_acc(c33));

endmodule
