`timescale 1ns/1ps

// ============================================================
// NMS STEP11 -- correctness testbench for weight_prefetch_engine.v.
// Real memory model (sim_word_mem, 2-cycle latency, same convention
// as tb_nms_dataflow_core.v/STEP8), a real weight SRAM
// (nms_weight_packed.v, N_SLOTS=1 for this isolated test), bit-exact
// data checking against a known fill pattern.
//
// Covers: n_tiles in {0,1,2,PFD,PFD+1,MAX_TILES-1,MAX_TILES}, PFD in
// {1,2,4,8}, back-to-back jobs (reset-between-jobs via job_active
// falling/rising, no explicit reset pulse), and a real memory-latency
// variant (sim_word_mem with an injected extra wait) to confirm the
// engine never double-fetches or overwrites a not-yet-committed tile.
// ============================================================
module sim_word_mem #(
    parameter ADDR_WIDTH = 23,
    parameter DEPTH      = 4096,
    parameter EXTRA_WAIT = 0
)(
    input  wire clk, rst,
    input  wire req, wr,
    input  wire [ADDR_WIDTH-1:0] addr,
    input  wire [15:0] wdata,
    input  wire lb_n, ub_n,
    output reg [15:0] rdata,
    output reg ready
);
    reg [15:0] mem [0:DEPTH-1];
    reg [3:0] state;
    reg [ADDR_WIDTH-1:0] addr_reg;
    reg [7:0] wait_cnt;
    localparam ST_IDLE=0, ST_WAIT=1, ST_EXTRA=2;
    always @(posedge clk) begin
        if (rst) begin state<=ST_IDLE; ready<=0; rdata<=0; end
        else begin
            ready <= 0;
            case (state)
                ST_IDLE: if (req) begin
                    addr_reg <= addr;
                    if (wr) begin
                        if (!lb_n) mem[addr][7:0] <= wdata[7:0];
                        if (!ub_n) mem[addr][15:8] <= wdata[15:8];
                    end
                    wait_cnt <= EXTRA_WAIT[7:0];
                    state <= ST_WAIT;
                end
                ST_WAIT: begin
                    if (wait_cnt != 0) begin
                        wait_cnt <= wait_cnt - 1'b1;
                    end else begin
                        rdata <= mem[addr_reg];
                        ready <= 1;
                        state <= ST_IDLE;
                    end
                end
                default: state <= ST_IDLE;
            endcase
        end
    end
endmodule

module tb;
    parameter ADDR_WIDTH = 23;
    parameter DATA_WIDTH = 8;
    parameter P_IN = 8;
    parameter MAX_TILES = 16;
    parameter PFD = 4;
    parameter EXTRA_WAIT = 0;
    localparam TIW = $clog2(MAX_TILES);
    localparam CNTW = $clog2(MAX_TILES+1);

    reg clk = 0;
    always #5 clk = ~clk;
    reg rst;

    reg job_active;
    reg [ADDR_WIDTH-1:0] w_base;
    reg [15:0] n_tiles;
    reg [CNTW-1:0] consumed_count;

    wire wgt_fill_we;
    wire [TIW-1:0] wgt_fill_addr;
    wire [DATA_WIDTH*P_IN-1:0] wgt_fill_data;
    wire [CNTW-1:0] ready_count;

    wire mem_req, mem_wr;
    wire [ADDR_WIDTH-1:0] mem_addr;
    wire [15:0] mem_wdata, mem_rdata;
    wire mem_lb_n, mem_ub_n, mem_ready;

    weight_prefetch_engine #(
        .DATA_WIDTH(DATA_WIDTH), .P_IN(P_IN), .ADDR_WIDTH(ADDR_WIDTH),
        .MAX_TILES(MAX_TILES), .PREFETCH_DISTANCE(PFD)
    ) dut (
        .clk(clk), .rst(rst),
        .job_active(job_active), .w_base(w_base), .n_tiles(n_tiles),
        .consumed_count(consumed_count),
        .wgt_fill_we(wgt_fill_we), .wgt_fill_addr(wgt_fill_addr), .wgt_fill_data(wgt_fill_data),
        .ready_count(ready_count),
        .mem_req(mem_req), .mem_wr(mem_wr), .mem_addr(mem_addr), .mem_wdata(mem_wdata),
        .mem_lb_n(mem_lb_n), .mem_ub_n(mem_ub_n), .mem_rdata(mem_rdata), .mem_ready(mem_ready)
    );

    sim_word_mem #(.ADDR_WIDTH(ADDR_WIDTH), .DEPTH(4096), .EXTRA_WAIT(EXTRA_WAIT)) u_mem (
        .clk(clk), .rst(rst), .req(mem_req), .wr(mem_wr), .addr(mem_addr), .wdata(mem_wdata),
        .lb_n(mem_lb_n), .ub_n(mem_ub_n), .rdata(mem_rdata), .ready(mem_ready)
    );

    // real weight SRAM (N_SLOTS=1) -- same module as production
    reg  wgt_rd_en;
    reg  [TIW-1:0] wgt_rd_addr;
    wire signed [DATA_WIDTH*P_IN-1:0] wgt_rd_data;
    nms_weight_packed #(.DATA_WIDTH(DATA_WIDTH), .P_IN(P_IN), .N_SLOTS(1), .MAX_TILES(MAX_TILES)) u_sram (
        .clk(clk), .rst(rst),
        .fill_we(wgt_fill_we), .fill_addr_flat(wgt_fill_addr), .fill_data_flat(wgt_fill_data),
        .rd_en(wgt_rd_en), .rd_addr_flat(wgt_rd_addr), .rd_data_flat(wgt_rd_data)
    );

    task automatic poke_word(input [ADDR_WIDTH-1:0] byte_addr, input [7:0] val);
        reg [ADDR_WIDTH-2:0] wa;
        begin
            wa = byte_addr[ADDR_WIDTH-1:1];
            if (byte_addr[0]==1'b0) u_mem.mem[wa][7:0] = val; else u_mem.mem[wa][15:8] = val;
        end
    endtask

    integer errors, tests;

    // fill PSRAM with a distinct pattern per tile: tile t, lane k -> (t*8+k) mod 251 (avoid trivial repeats)
    task automatic fill_pattern(input [ADDR_WIDTH-1:0] base, input integer count);
        integer t, k;
        begin
            for (t = 0; t < count; t = t + 1)
                for (k = 0; k < P_IN; k = k + 1)
                    poke_word(base + t*P_IN + k, (t*8+k) % 251);
        end
    endtask

    // simulated consumer: follows ready_count with a small lag (models
    // a real nms_memory_manager.v consuming tiles about as fast as
    // they arrive) -- driven every cycle while a job is active, unless
    // freeze_consumer is set (dedicated windowing-cap test below).
    reg freeze_consumer;
    always @(posedge clk) begin
        if (rst || !job_active) consumed_count <= {CNTW{1'b0}};
        else if (!freeze_consumer && consumed_count < ready_count) consumed_count <= consumed_count + 1'b1;
    end

    task automatic run_job(input [ADDR_WIDTH-1:0] base, input integer count, input integer watchdog);
        integer wd, t, k;
        reg [7:0] expected;
        begin
            w_base = base; n_tiles = count[15:0];
            job_active = 1'b1;
            wd = 0;
            while (ready_count < count[CNTW-1:0] && wd < watchdog) begin @(posedge clk); wd = wd + 1; end
            @(posedge clk); #1;
            tests = tests + 1;
            if (ready_count !== count[CNTW-1:0]) begin
                $display("FAIL n_tiles=%0d PFD=%0d: ready_count=%0d expected=%0d (watchdog=%0d)", count, PFD, ready_count, count, wd);
                errors = errors + 1;
            end else begin
                // verify every tile's data bit-exact
                for (t = 0; t < count; t = t + 1) begin
                    wgt_rd_addr = t[TIW-1:0]; wgt_rd_en = 1'b1;
                    @(posedge clk); @(posedge clk); #1; // 2-cycle SRAM read latency
                    for (k = 0; k < P_IN; k = k + 1) begin
                        expected = (t*8+k) % 251;
                        if (wgt_rd_data[k*DATA_WIDTH +: DATA_WIDTH] !== expected) begin
                            $display("FAIL n_tiles=%0d PFD=%0d tile=%0d lane=%0d: got=%0d expected=%0d",
                                count, PFD, t, k, wgt_rd_data[k*DATA_WIDTH +: DATA_WIDTH], expected);
                            errors = errors + 1;
                        end
                    end
                end
                $display("PASS n_tiles=%0d PFD=%0d EXTRA_WAIT=%0d: ready_count=%0d, all tiles bit-exact (cycles=%0d)",
                    count, PFD, EXTRA_WAIT, ready_count, wd);
            end
            job_active = 1'b0;
            repeat(3) @(posedge clk);
        end
    endtask

    initial begin
        errors = 0; tests = 0;
        rst = 1; job_active = 0; w_base = 0; n_tiles = 0; consumed_count = 0; wgt_rd_en = 0; wgt_rd_addr = 0; freeze_consumer = 0;
        repeat(4) @(posedge clk);
        rst = 0;
        @(posedge clk);

        fill_pattern(23'h1000, MAX_TILES);

        // edge cases explicitly requested: 0,1,2,PFD,PFD+1,MAX_TILES-1,MAX_TILES
        // (PFD/PFD+1 only make sense as a job size when they fit inside
        // MAX_TILES -- the weight SRAM itself only holds MAX_TILES
        // tiles, so PFD>=MAX_TILES's own "PFD,PFD+1" cases are simply
        // n/a, not a real scenario to test; MAX_TILES-1/MAX_TILES below
        // already cover the near-/at-capacity edge regardless of PFD)
        run_job(23'h1000, 0, 200);
        run_job(23'h1000, 1, 200);
        run_job(23'h1000, 2, 200);
        if (PFD < MAX_TILES) begin
            run_job(23'h1000, PFD, 500);
            run_job(23'h1000, PFD+1, 500);
        end
        run_job(23'h1000, MAX_TILES-1, 2000);
        run_job(23'h1000, MAX_TILES, 2000);

        // back-to-back jobs, no explicit reset between them (job_active
        // falling then rising, exercising the "!job_active" reset path)
        run_job(23'h1000, 3, 500);
        run_job(23'h1000, 5, 500);

        // ERR-0015 regression: PREFETCH_DISTANCE >= 2**CNTW (here 32,
        // CNTW=5 bits for MAX_TILES=16) must NOT silently truncate to 0
        // and deadlock (window_limit==consumed_count forever). Only
        // meaningful when PFD is actually >= 32; skipped otherwise so
        // this file stays valid across all PFD values it's compiled
        // with (per the sweep in EXP-0023/EXP-0024).
        if (PFD >= 32) begin : err0015_large_pfd_test
            run_job(23'h1000, MAX_TILES, 2000);
        end

        // dedicated windowing-cap test: consumer NEVER consumes
        // (freeze_consumer=1, consumed_count stuck at 0) -- ready_count
        // must stop advancing at exactly min(PFD,MAX_TILES) tiles (the
        // n_tiles/MAX_TILES ceiling binds first whenever PFD>=MAX_TILES,
        // e.g. the ERR-0015 large-PFD case above), never fetching
        // further ahead than whichever bound applies, and must NEVER
        // re-fetch/overwrite once frozen (checked by waiting well past
        // when an unbounded engine would have finished all MAX_TILES,
        // then confirming ready_count is still exactly at that bound).
        begin : window_cap_test
            integer wd2;
            integer expected_cap;
            tests = tests + 1;
            expected_cap = (PFD < MAX_TILES) ? PFD : MAX_TILES;
            freeze_consumer = 1'b1;
            w_base = 23'h1000; n_tiles = MAX_TILES[15:0];
            job_active = 1'b1;
            for (wd2 = 0; wd2 < 500; wd2 = wd2 + 1) @(posedge clk);
            if (ready_count !== expected_cap[CNTW-1:0]) begin
                $display("FAIL windowing cap: ready_count=%0d expected exactly %0d (PFD=%0d, MAX_TILES=%0d, consumer frozen at 0)", ready_count, expected_cap, PFD, MAX_TILES);
                errors = errors + 1;
            end else
                $display("PASS windowing cap: ready_count correctly capped at %0d (PFD=%0d, MAX_TILES=%0d) with consumer frozen", expected_cap, PFD, MAX_TILES);
            job_active = 1'b0;
            freeze_consumer = 1'b0;
            repeat(3) @(posedge clk);
        end

        $display("=== %0d/%0d tests, %0d errors (PFD=%0d, MAX_TILES=%0d, EXTRA_WAIT=%0d) ===",
            tests-errors, tests, errors, PFD, MAX_TILES, EXTRA_WAIT);
        if (errors == 0) $display("ALL TESTS PASSED (tb_weight_prefetch, PFD=%0d, EXTRA_WAIT=%0d)", PFD, EXTRA_WAIT);
        $finish;
    end
endmodule