// IM2COL + GEMM Convolution Engine
module im2col_gemm #(parameter DW=16, KH=3, ICH=16, OCH=16, OW=28, OH=28)(
    input clk, rst_n, input en, input start,
    // Input feature map
    input signed [DW-1:0] ifm [0:ICH*32*32-1], // Simplified flat array
    // Weight matrix (already reshaped to OCH × ICH*KH*KH)
    input signed [DW-1:0] wt_mat [0:OCH*ICH*KH*KH-1],
    // Output
    output reg signed [DW-1:0] ofm_out, output reg ofm_valid,
    output reg [7:0] out_ch, output reg [9:0] out_px,
    output reg done
);
    // IM2COL: generate unfolded column on the fly
    reg signed [DW-1:0] col_buf [0:ICH*KH*KH-1]; // One column of im2col matrix
    reg [7:0] oc_cnt, ic_cnt, kr_cnt, kc_cnt;
    reg [9:0] px_cnt;
    reg [3:0] state;
    reg signed [DW*2-1:0] psum;
    // IM2COL address generation
    wire [9:0] base_y = (px_cnt / OW) * 1; // Stride=1
    wire [9:0] base_x = px_cnt % OW;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin state<=0; psum<='0; ofm_out<='0; ofm_valid<=0; done<=0;
            oc_cnt<=0; ic_cnt<=0; kr_cnt<=0; kc_cnt<=0; px_cnt<=0; out_ch<=0; out_px<=0; end
        else case(state)
          0: if(start) begin state<=1; px_cnt<=0; oc_cnt<=0; ic_cnt<=0; kr_cnt<=0; kc_cnt<=0; psum<='0; end
          1: begin // IM2COL + GEMM combined
                // Compute one MAC: col_buf element * weight element
                psum <= psum + col_buf[ic_cnt*KH*KH+kr_cnt*KH+kc_cnt] *
                        wt_mat[oc_cnt*ICH*KH*KH+ic_cnt*KH*KH+kr_cnt*KH+kc_cnt];
                kc_cnt<=kc_cnt+1;
                if(kc_cnt>=KH-1) begin kc_cnt<=0; kr_cnt<=kr_cnt+1;
                  if(kr_cnt>=KH-1) begin kr_cnt<=0; ic_cnt<=ic_cnt+1;
                    if(ic_cnt>=ICH-1) begin ic_cnt<=0;
                        ofm_out<=psum[DW*2-1:DW]; ofm_valid<=1;
                        out_ch<=oc_cnt; out_px<=px_cnt;
                        psum<='0; oc_cnt<=oc_cnt+1;
                        if(oc_cnt>=OCH-1) begin oc_cnt<=0; px_cnt<=px_cnt+1;
                            if(px_cnt>=OW*OH-1) state<=3;
                        end
                    end
                  end
                end
            end
          3: done<=1;
        endcase
    end
endmodule