// Weight Pruning Engine — Structured + Unstructured pruning
module pruning_engine #(parameter DW=16, NW=256, THRESH_W=8)(
    input clk, rst_n, input en,
    input [1:0] pr_mode, // 00=none, 01=unstructured, 10=structured_row, 11=structured_col
    input signed [THRESH_W-1:0] threshold,
    input signed [DW-1:0] weight_in, input [$clog2(NW)-1:0] w_idx, input w_valid,
    output reg signed [DW-1:0] weight_out, output reg w_out_valid, output reg is_zero,
    // Structured pruning control
    input [$clog2(NW)-1:0] row_len, // Elements per row for structured pruning
    output reg [15:0] sparsity,     // Current sparsity percentage
    output reg done
);
    reg [15:0] total_cnt, zero_cnt;
    reg signed [THRESH_W-1:0] row_max; // For structured: track row max
    reg [$clog2(NW)-1:0] row_pos;
    // Unstructured: |w| < threshold => zero
    wire unstruct_zero = (weight_in >= 0 ? weight_in : -weight_in) < threshold;
    // Structured row: track max in row, prune entire row if max < threshold
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin weight_out<='0; w_out_valid<=0; is_zero<=0; total_cnt<=0; zero_cnt<=0; sparsity<=0; done<=0; row_max<=0; row_pos<=0; end
        else if(en && w_valid) begin
            total_cnt<=total_cnt+1; w_out_valid<=1;
            case(pr_mode)
              2'd01: begin // Unstructured
                is_zero<=unstruct_zero; weight_out<=unstruct_zero?'0:weight_in;
                if(unstruct_zero) zero_cnt<=zero_cnt+1;
              end
              2'd10: begin // Structured row
                row_pos<=row_pos+1;
                if(weight_in>row_max||(weight_in<0&&-weight_in>row_max)) row_max<=weight_in>=0?weight_in:-weight_in;
                if(row_pos>=row_len-1) begin // End of row
                    if(row_max<threshold) begin is_zero<=1; zero_cnt<=zero_cnt+row_len; weight_out<='0; end
                    else begin is_zero<=0; weight_out<=weight_in; end
                    row_max<=0; row_pos<=0;
                end else begin is_zero<=0; weight_out<=weight_in; end
              end
              default: begin weight_out<=weight_in; is_zero<=0; end
            endcase
            sparsity<=(zero_cnt*100)/total_cnt;
            if(total_cnt>=NW) done<=1;
        end else begin w_out_valid<=0; done<=0; end
    end
endmodule