`timescale 1ns / 1ps

module popcount_pipeline #(
    parameter int WIDTH = 512,
    parameter int ROW_BITS = 12,
    parameter int OUT_WIDTH = $clog2(WIDTH + 1)
) (
    input  logic                 clk,
    input  logic                 rst_n,
    input  logic                 valid_i,
    input  logic [ROW_BITS-1:0]  row_i,
    input  logic [WIDTH-1:0]     bits_i,
    output logic                 valid_o,
    output logic [ROW_BITS-1:0]  row_o,
    output logic [OUT_WIDTH-1:0] count_o
);
    localparam int GROUP = 8;
    localparam int NUM_GROUPS = WIDTH / GROUP;
    localparam int GROUP_COUNT_WIDTH = $clog2(GROUP + 1);
    localparam int STAGE1_GROUPS = NUM_GROUPS / 4;

    logic [GROUP_COUNT_WIDTH-1:0] group_counts [0:NUM_GROUPS-1];
    logic [OUT_WIDTH-1:0] partial_q [0:3];
    logic [OUT_WIDTH-1:0] sum_q;
    logic [ROW_BITS-1:0] row_s1_q, row_s2_q;
    logic valid_s1_q, valid_s2_q;
    logic [OUT_WIDTH-1:0] partial_comb [0:3];

    always_comb begin
        for (int g = 0; g < NUM_GROUPS; g++) begin
            group_counts[g] = '0;
            for (int b = 0; b < GROUP; b++) begin
                group_counts[g] = group_counts[g] + bits_i[g*GROUP + b];
            end
        end

        for (int p = 0; p < 4; p++) begin
            partial_comb[p] = '0;
            for (int g = 0; g < STAGE1_GROUPS; g++) begin
                partial_comb[p] = partial_comb[p] + group_counts[p*STAGE1_GROUPS + g];
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            valid_s1_q <= 1'b0;
            valid_s2_q <= 1'b0;
            valid_o    <= 1'b0;
            row_s1_q   <= '0;
            row_s2_q   <= '0;
            row_o      <= '0;
            sum_q      <= '0;
            count_o    <= '0;
            for (int p = 0; p < 4; p++) partial_q[p] <= '0;
        end else begin
            valid_s1_q <= valid_i;
            row_s1_q   <= row_i;
            for (int p = 0; p < 4; p++) begin
                partial_q[p] <= partial_comb[p];
            end

            valid_s2_q <= valid_s1_q;
            row_s2_q   <= row_s1_q;
            sum_q      <= partial_q[0] + partial_q[1] + partial_q[2] + partial_q[3];

            valid_o <= valid_s2_q;
            row_o   <= row_s2_q;
            count_o <= sum_q;
        end
    end
endmodule