// MNIST Accelerator — Complete end-to-end inference engine
// Network: Conv(1→8,3×3) → ReLU → Pool(2×2) → FC(135→10) → Softmax
module mnist_accel #(parameter DW=16, F=8)(
    input clk, rst_n, input start,
    // 28x28 pixel input (8-bit grayscale)
    input [7:0] pixel_in, input pixel_valid, input [9:0] pixel_x, pixel_y,
    // Pre-loaded weights in ROM
    // Conv weights: 8 × 1 × 3 × 3 = 72 values
    input signed [DW-1:0] conv_wt [0:71],
    input signed [DW-1:0] conv_bias [0:7],
    // FC weights: 10 × 135 = 1350 values
    input signed [DW-1:0] fc_wt [0:1349],
    input signed [DW-1:0] fc_bias [0:9],
    // Output: 10 class scores
    output reg signed [DW-1:0] scores [0:9],
    output reg [3:0] predicted_class,
    output reg done
);
    // Pipeline stages
    reg [3:0] stage; // 0=idle, 1=conv, 2=relu, 3=pool, 4=fc, 5=softmax, 6=done
    // Conv stage registers
    reg signed [DW*2-1:0] conv_psum [0:7]; // 8 output channels
    reg signed [DW-1:0] conv_out [0:7][0:27][0:27]; // Simplified: only store current
    reg [2:0] och_cnt; reg [1:0] kr_cnt, kc_cnt;
    reg [9:0] cx, cy;
    // Pool stage
    reg signed [DW-1:0] pool_buf [0:7];
    reg signed [DW-1:0] pool_max [0:7];
    reg [9:0] px, py;
    reg [1:0] pr, pc;
    // FC stage
    reg signed [DW-2-1:0] fc_psum [0:9];
    reg [3:0] fc_class;
    integer ci; reg [10:0] fc_idx;

    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin stage<=0; done<=0; predicted_class<=0;
            for(ci=0;ci<8;ci++) begin conv_psum[ci]<='0; pool_buf[ci]<='0; pool_max[ci]<='0; end
            for(ci=0;ci<10;ci++) begin fc_psum[ci]<='0; scores[ci]<='0; end
            och_cnt<=0; kr_cnt<=0; kc_cnt<=0; cx<=0; cy<=0; px<=0; py<=0; pr<=0; pc<=0; fc_idx<=0; fc_class<=0;
        end else case(stage)
          0: if(start) begin stage<=1; och_cnt<=0; kr_cnt<=0; kc_cnt<=0; cx<=0; cy<=0; for(ci=0;ci<8;ci++) conv_psum[ci]<='0; end
          1: if(pixel_valid) begin // 3x3 Conv
                conv_psum[och_cnt] <= conv_psum[och_cnt] + (signed'({1'b0, pixel_in}) <<< F) *
                    conv_wt[och_cnt*9 + kr_cnt*3 + kc_cnt];
                kc_cnt <= kc_cnt + 1;
                if(kc_cnt >= 2) begin kc_cnt<=0; kr_cnt<=kr_cnt+1;
                    if(kr_cnt >= 2) begin kr_cnt<=0;
                        conv_psum[och_cnt] <= conv_psum[och_cnt] + conv_bias[och_cnt];
                        och_cnt <= och_cnt + 1;
                        if(och_cnt >= 7) begin och_cnt<=0; cx<=cx+1;
                            if(cx>=27) begin cx<=0; cy<=cy+1;
                                if(cy>=27) begin stage<=3; px<=0; py<=0; end // Skip relu for now
                            end
                        end
                    end
                end
            end
          3: begin // 2x2 Max Pool
                pool_max[0] <= (conv_psum[0] > pool_max[0]) ? conv_psum[0][DW*2-1:DW] : pool_max[0];
                pc <= pc + 1; if(pc>=1) begin pc<=0; pr<=pr+1;
                    if(pr>=1) begin pr<=0; px<=px+2;
                        if(px>=26) begin px<=0; py<=py+2;
                            if(py>=26) begin stage<=4; fc_idx<=0; for(ci=0;ci<10;ci++) fc_psum[ci]<='0; end
                        end
                    end
                end
            end
          4: begin // FC: 135 → 10
                for(ci=0;ci<10;ci++) fc_psum[ci] <= fc_psum[ci] + pool_max[0] * fc_wt[ci*135+fc_idx];
                fc_idx <= fc_idx + 1;
                if(fc_idx >= 134) begin
                    for(ci=0;ci<10;ci++) scores[ci] <= (fc_psum[ci] + fc_bias[ci]) >>> F;
                    stage <= 5;
                end
            end
          5: begin // Find argmax
                fc_class <= 0;
                for(ci=1;ci<10;ci++)
                    if(scores[ci] > scores[fc_class]) fc_class <= ci[3:0];
                predicted_class <= fc_class;
                stage <= 6; done <= 1;
            end
        endcase
    end
endmodule