// Inference Scheduler — Layer-by-layer execution with memory management
module infer_scheduler #(parameter DW=16, MAX_LAYERS=50, MAX_BUF=8, AW=16)(
    input clk, rst_n, input start,
    input [7:0] num_layers,
    // Layer configuration
    input [7:0] layer_type [0:MAX_LAYERS-1],  // 0=conv, 1=fc, 2=pool, 3=act, 4=concat
    input [AW-1:0] layer_ifm_size [0:MAX_LAYERS-1],
    input [AW-1:0] layer_ofm_size [0:MAX_LAYERS-1],
    input [AW-1:0] layer_wt_size  [0:MAX_LAYERS-1],
    // Buffer allocation
    output reg [2:0] ifm_buf_id [0:MAX_LAYERS-1],
    output reg [2:0] ofm_buf_id [0:MAX_LAYERS-1],
    // Compute engine control
    output reg [7:0] cur_layer, output reg layer_start, input layer_done,
    output reg [1:0] exec_cmd, // 00=conv, 01=fc, 10=pool, 11=act
    output reg done
);
    reg [7:0] layer_cnt;
    reg [3:0] state; // 0=idle, 1=alloc, 2=exec, 3=next, 4=done
    reg [2:0] next_buf;
    reg [2:0] buf_used [0:MAX_BUF-1];
    integer bi;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin state<=0; layer_cnt<=0; next_buf<=0; done<=0; cur_layer<=0; layer_start<=0; exec_cmd<=0;
            for(bi=0;bi<MAX_BUF;bi++) buf_used[bi]<=0; for(bi=0;bi<MAX_LAYERS;bi++) begin ifm_buf_id[bi]<=0; ofm_buf_id[bi]<=0; end end
        else case(state)
          0: if(start) begin state<=1; layer_cnt<=0; next_buf<=0; end
          1: begin // Allocate buffers for current layer
                ifm_buf_id[layer_cnt] <= (layer_cnt==0) ? 0 : ofm_buf_id[layer_cnt-1];
                ofm_buf_id[layer_cnt] <= next_buf;
                buf_used[next_buf] <= 1;
                next_buf <= next_buf + 1;
                if(next_buf >= MAX_BUF-1) next_buf <= 0; // Wrap around
                state <= 2; layer_start <= 1;
            end
          2: begin layer_start<=0;
                if(layer_done) begin state<=3;
                    // Free input buffer if no longer needed
                    buf_used[ofm_buf_id[layer_cnt-1]] <= 0;
                end
            end
          3: begin
                layer_cnt <= layer_cnt + 1;
                cur_layer <= layer_cnt + 1;
                if(layer_cnt >= num_layers - 1) state <= 4;
                else state <= 1;
            end
          4: done <= 1;
        endcase
    end
endmodule