// CNN Inference Engine — Configurable multi-layer CNN
module cnn_inference #(parameter DW=16, MAX_LAYERS=20, MAX_CH=64, MAX_DIM=128)(
    input clk, rst_n, input start,
    input [7:0] num_layers,
    // Layer config
    input [7:0] layer_och [0:MAX_LAYERS-1],
    input [7:0] layer_ich [0:MAX_LAYERS-1],
    input [2:0] layer_ksz [0:MAX_LAYERS-1], // 1,3,5,7
    input       layer_has_relu [0:MAX_LAYERS-1],
    input       layer_has_pool [0:MAX_LAYERS-1],
    input [7:0] layer_stride [0:MAX_LAYERS-1],
    // Feature map buffer (ping-pong)
    input signed [DW-1:0] fm_buf_a [0:MAX_CH*MAX_DIM*MAX_DIM-1],
    input signed [DW-1:0] fm_buf_b [0:MAX_CH*MAX_DIM*MAX_DIM-1],
    // Weight ROM
    input signed [DW-1:0] weight_rom [0:65535],
    input signed [DW-1:0] bias_rom [0:MAX_CH-1],
    // Output
    output reg signed [DW-1:0] final_out, output reg final_valid,
    output reg [7:0] out_ch, output reg [14:0] out_addr,
    output reg done
);
    reg [7:0] layer_cnt, och, ich, kr, kc;
    reg [14:0] ox, oy, wt_offset;
    reg signed [DW*2-1:0] psum;
    reg [3:0] state;
    reg using_buf_a; // Ping-pong flag
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin state<=0; layer_cnt<=0; och<=0; ich<=0; kr<=0; kc<=0;
            ox<=0; oy<=0; psum<='0; wt_offset<=0; using_buf_a<=1; done<=0;
            final_out<='0; final_valid<=0; out_ch<=0; out_addr<=0; end
        else case(state)
          0: if(start) begin state<=1; layer_cnt<=0; wt_offset<=0; using_buf_a<=1; end
          1: begin // Conv compute
                psum <= psum + fm_buf_a[ich*MAX_DIM*MAX_DIM + (oy+kr)*MAX_DIM + ox+kc] *
                    weight_rom[wt_offset + och*layer_ich[layer_cnt]*layer_ksz[layer_cnt]*layer_ksz[layer_cnt] + ich*layer_ksz[layer_cnt]*layer_ksz[layer_cnt] + kr*layer_ksz[layer_cnt] + kc];
                kc <= kc + 1;
                if(kc >= layer_ksz[layer_cnt]-1) begin kc<=0; kr<=kr+1;
                  if(kr >= layer_ksz[layer_cnt]-1) begin kr<=0; ich<=ich+1;
                    if(ich >= layer_ich[layer_cnt]-1) begin ich<=0;
                        psum <= psum + bias_rom[och]; // Add bias
                        // ReLU
                        if(layer_has_relu[layer_cnt] && psum[DW*2-1]) psum <= '0;
                        final_out <= psum[DW*2-1:DW]; final_valid <= 1;
                        out_ch <= och; out_addr <= och*MAX_DIM*MAX_DIM + oy*MAX_DIM + ox;
                        psum <= '0; och <= och + 1;
                        if(och >= layer_och[layer_cnt]-1) begin och<=0; ox<=ox+layer_stride[layer_cnt];
                          if(ox >= MAX_DIM-1) begin ox<=0; oy<=oy+layer_stride[layer_cnt];
                            if(oy >= MAX_DIM-1) begin
                                wt_offset <= wt_offset + layer_och[layer_cnt]*layer_ich[layer_cnt]*layer_ksz[layer_cnt]*layer_ksz[layer_cnt];
                                layer_cnt <= layer_cnt + 1;
                                using_buf_a <= ~using_buf_a;
                                if(layer_cnt >= num_layers-1) state <= 2;
                            end
                          end
                        end
                    end
                  end
                end
            end
          2: done <= 1;
        endcase
    end
endmodule