`timescale 1ns/1ps

// ============================================================
// NMS STEP15 (continuation) -- bit-exact + timing regression for
// psram_controller_dual32.v: real weight_prefetch_engine_wide.v
// (STEP14, UNMODIFIED) at MEM_DATA_WIDTH=32, driving the REAL dual-
// chip 32-bit controller against TWO real psram_model.v instances,
// with the real production nms_weight_packed.v SRAM as the fill
// target -- same bit-exact methodology as tb_weight_prefetch_wide.v
// (STEP14), byte-level pattern (t*8+k)%251, per-tile read-back
// verification.
//
// Covers: bit-exact data (lane ordering: chip0=low16/chip1=high16 of
// each 32-bit word), page-hit/open/close/boundary behavior (inherited
// unmodified from the real psram_controller.v, exercised identically
// per physical chip), reset behavior, back-to-back transactions
// (n_tiles edge cases), and real cycles/tile timing (expected ~9).
// ============================================================
module tb;
    parameter ADDR_WIDTH = 23;
    parameter DATA_WIDTH = 8;
    parameter P_IN = 8;
    parameter MAX_TILES = 512;
    parameter PFD = 600;
    localparam TIW = $clog2(MAX_TILES);
    localparam CNTW = $clog2(MAX_TILES+1);
    localparam CLK_PERIOD = 12.5; // 80MHz, matches every real-PSRAM benchmark in this project

    reg clk = 0;
    always #(CLK_PERIOD/2.0) 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;
    wire [ADDR_WIDTH-1:0] mem_addr;
    wire [31:0] mem_rdata;
    wire mem_ready;

    weight_prefetch_engine_wide #(
        .DATA_WIDTH(DATA_WIDTH), .P_IN(P_IN), .ADDR_WIDTH(ADDR_WIDTH),
        .MAX_TILES(MAX_TILES), .PREFETCH_DISTANCE(PFD), .MEM_DATA_WIDTH(32)
    ) 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_addr(mem_addr), .mem_rdata(mem_rdata), .mem_ready(mem_ready)
    );

    // weight fetch never writes -- tie the real controller's write-
    // side inputs to constants (this engine has no write path, same
    // as its own real 16-bit counterpart weight_prefetch_engine.v)
    wire mem_wr = 1'b0;
    wire [31:0] mem_wdata = 32'h0;
    wire [1:0] mem_lb_n = 2'b00, mem_ub_n = 2'b00; // always both bytes of both chips

    wire [ADDR_WIDTH-1:0] p0_a, p1_a;
    wire [15:0] p0_dq, p1_dq;
    wire p0_ce_n, p0_oe_n, p0_we_n, p0_lb_n, p0_ub_n, p0_zz_n;
    wire p1_ce_n, p1_oe_n, p1_we_n, p1_lb_n, p1_ub_n, p1_zz_n;
    wire lane_sync_error;

    psram_controller_dual32 #(.ADDR_WIDTH(ADDR_WIDTH), .CLK_FREQ_MHZ(80)) u_dual (
        .clk(clk), .rst(rst),
        .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), .lane_sync_error(lane_sync_error),
        .psram0_a(p0_a), .psram0_dq(p0_dq),
        .psram0_ce_n(p0_ce_n), .psram0_oe_n(p0_oe_n), .psram0_we_n(p0_we_n),
        .psram0_lb_n(p0_lb_n), .psram0_ub_n(p0_ub_n), .psram0_zz_n(p0_zz_n),
        .psram1_a(p1_a), .psram1_dq(p1_dq),
        .psram1_ce_n(p1_ce_n), .psram1_oe_n(p1_oe_n), .psram1_we_n(p1_we_n),
        .psram1_lb_n(p1_lb_n), .psram1_ub_n(p1_ub_n), .psram1_zz_n(p1_zz_n)
    );

    psram_model #(.ADDR_WIDTH(ADDR_WIDTH), .DATA_WIDTH(16), .DEPTH(131072)) u_psram0 (
        .clk(clk), .a(p0_a), .dq(p0_dq),
        .ce_n(p0_ce_n), .oe_n(p0_oe_n), .we_n(p0_we_n),
        .lb_n(p0_lb_n), .ub_n(p0_ub_n), .zz_n(p0_zz_n)
    );
    psram_model #(.ADDR_WIDTH(ADDR_WIDTH), .DATA_WIDTH(16), .DEPTH(131072)) u_psram1 (
        .clk(clk), .a(p1_a), .dq(p1_dq),
        .ce_n(p1_ce_n), .oe_n(p1_oe_n), .we_n(p1_we_n),
        .lb_n(p1_lb_n), .ub_n(p1_ub_n), .zz_n(p1_zz_n)
    );

    // real production weight SRAM, N_SLOTS=1, fed by the fill port
    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)
    );

    // Poke the SAME byte-level pattern as tb_weight_prefetch_wide.v
    // (STEP14): tile t, lane k -> (t*8+k) % 251. weight_prefetch_
    // engine_wide.v's own mem_addr is a byte address; the real dual32
    // controller's own mem_addr is a WORD address (4 bytes/word at
    // 32-bit) -- poke directly into each chip's own byte-addressable
    // backing array via the SAME byte-address convention psram_model.v
    // itself uses elsewhere in this project (word_addr = byte_addr>>1
    // PER CHIP, since each chip is still a 16-bit device internally;
    // for the dual32 mapping, chip0 holds bits[15:0] of 32-bit word
    // W=byte_addr>>2, chip1 holds bits[31:16]).
    task automatic poke_byte(input [ADDR_WIDTH-1:0] byte_addr, input [7:0] val);
        reg [ADDR_WIDTH-1:0] word32_addr;
        reg [1:0] byte_in_word32;
        begin
            word32_addr    = byte_addr >> 2;
            byte_in_word32 = byte_addr[1:0];
            case (byte_in_word32)
                2'd0: u_psram0.mem[word32_addr][7:0]   = val;
                2'd1: u_psram0.mem[word32_addr][15:8]  = val;
                2'd2: u_psram1.mem[word32_addr][7:0]   = val;
                2'd3: u_psram1.mem[word32_addr][15:8]  = val;
            endcase
        end
    endtask

    integer errors, tests;

    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_byte(base + t*P_IN + k, (t*8+k) % 251);
        end
    endtask

    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

    integer cyc;
    always @(posedge clk) if (!rst) cyc <= cyc + 1;
    reg trace_on;
    always @(posedge clk) begin
        if (trace_on && mem_req)
            $display("  [%0d] mem_req addr=%0h (chip_word_addr=%0h) wr=%0d", cyc, mem_addr, u_dual.chip_word_addr, mem_wr);
        if (trace_on && mem_ready)
            $display("  [%0d] mem_ready rdata=%08h (rdata0=%04h rdata1=%04h)", cyc, mem_rdata, u_dual.rdata0, u_dual.rdata1);
        if (trace_on && wgt_fill_we)
            $display("  [%0d] wgt_fill_we addr=%0d data=%016h", cyc, wgt_fill_addr, wgt_fill_data);
    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: ready_count=%0d expected=%0d (watchdog=%0d)", count, ready_count, count, wd);
                errors = errors + 1;
            end else begin
                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;
                    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 tile=%0d lane=%0d: got=%0d expected=%0d",
                                count, t, k, wgt_rd_data[k*DATA_WIDTH +: DATA_WIDTH], expected);
                            errors = errors + 1;
                        end
                    end
                end
                if (lane_sync_error) begin
                    $display("FAIL n_tiles=%0d: lane_sync_error latched -- chips diverged", count);
                    errors = errors + 1;
                end
                $display("PASS n_tiles=%0d: ready_count=%0d, all tiles bit-exact, lane_sync_error=0 (cycles=%0d)",
                    count, ready_count, wd);
            end
            job_active = 1'b0;
            repeat(3) @(posedge clk);
        end
    endtask

    integer t0, t1;
    initial begin
        errors = 0; tests = 0; cyc = 0; freeze_consumer = 0;
        rst = 1; job_active = 0; w_base = 0; n_tiles = 0; consumed_count = 0; wgt_rd_en = 0; wgt_rd_addr = 0; trace_on = 0;
        repeat(5) @(posedge clk);
        rst = 0;

        // real ~150us power-up wait, both chips (matches every other
        // real-PSRAM testbench in this project) -- forgetting this
        // was an earlier test-setup bug in this file (all requests
        // issued during STATE_INIT/STATE_CR_INIT are simply never
        // accepted), not an RTL defect.
        wait (u_dual.u_ctrl0.state == u_dual.u_ctrl0.STATE_IDLE);
        wait (u_dual.u_ctrl1.state == u_dual.u_ctrl1.STATE_IDLE);
        @(posedge clk);

        fill_pattern(23'h60000, MAX_TILES);
        trace_on = 1'b1;

        // edge cases (STEP11-14 convention): 0,1,2,MAX_TILES-1,MAX_TILES,
        // back-to-back jobs (reset-free), matching prior discipline --
        // "counter-width bug at value 16" class explicitly re-tested here.
        run_job(23'h60000, 0, 500);
        run_job(23'h60000, 1, 500);
        run_job(23'h60000, 2, 500);
        run_job(23'h60000, 15, 2000);
        run_job(23'h60000, 16, 2000);
        run_job(23'h60000, MAX_TILES-1, 8000);

        // timed run for the real cycles/tile measurement (unconstrained
        // consumer, freeze_consumer=1, so ready_count races ahead as
        // fast as the real dual-chip physical interface allows)
        freeze_consumer = 1'b1;
        w_base = 23'h60000; n_tiles = MAX_TILES[15:0];
        job_active = 1'b1;
        t0 = cyc;
        wait (ready_count == MAX_TILES[CNTW-1:0]);
        t1 = cyc;
        $display("REAL DUAL-CHIP 32-BIT TIMING: MAX_TILES=%0d total_cycles=%0d cycles/tile=%0.4f lane_sync_error=%0d",
            MAX_TILES, t1-t0, (t1-t0)/(1.0*MAX_TILES), lane_sync_error);
        job_active = 1'b0;
        freeze_consumer = 1'b0;

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