`timescale 1ns/1ps

// ================================================================
// PSRAM PAGE MODE TEST
//
// Exercises the page-mode burst-read path added to
// rtl/psram_controller.v: the configuration-register load at
// power-up, fast same-page read continuations (tAPA instead of
// tAA) including byte-enable changes between words (the pattern
// int8_memory_access.v actually produces -- LB#/UB# alternate on
// nearly every access), page-boundary crossing within an open
// session, and the one condition that closes the page (a WRITE).
//
// sim/psram_model.v enforces real datasheet timing ($fatal on any
// violation), including the new page-mode tAPA/tAA continuation
// check -- so a passing run here is a real proof that the RTL
// waits the correct number of cycles, not just that data compares
// equal.
// ================================================================

module tb;

    localparam ADDR_WIDTH = 23;
    localparam DATA_WIDTH = 16;
    localparam CLK_FREQ_MHZ = 80;
    localparam CLK_PERIOD = 12.5; // 80 MHz

    reg clk;
    reg rst;

    reg                   mem_req;
    reg                   mem_wr;
    reg  [ADDR_WIDTH-1:0] mem_addr;
    reg  [DATA_WIDTH-1:0] mem_wdata;
    reg                   mem_lb_n;
    reg                   mem_ub_n;

    wire [DATA_WIDTH-1:0] mem_rdata;
    wire                  mem_ready;

    wire [ADDR_WIDTH-1:0] psram_a;
    wire [DATA_WIDTH-1:0] psram_dq;

    wire psram_ce_n;
    wire psram_oe_n;
    wire psram_we_n;
    wire psram_lb_n;
    wire psram_ub_n;
    wire psram_zz_n;

    // ============================================================
    // Expected cycle counts (must match the RTL's own formulas)
    // ============================================================

    localparam integer ACCESS_CYCLES = ((70 * CLK_FREQ_MHZ) + 999) / 1000;
    localparam integer PAGE_CYCLES   = ((20 * CLK_FREQ_MHZ) + 999) / 1000;

    // ============================================================
    // DUT
    // ============================================================

    psram_controller #(
        .ADDR_WIDTH   (ADDR_WIDTH),
        .DATA_WIDTH   (DATA_WIDTH),
        .CLK_FREQ_MHZ (CLK_FREQ_MHZ)
    ) dut (
        .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),

        .psram_a   (psram_a),
        .psram_dq  (psram_dq),

        .psram_ce_n(psram_ce_n),
        .psram_oe_n(psram_oe_n),
        .psram_we_n(psram_we_n),
        .psram_lb_n(psram_lb_n),
        .psram_ub_n(psram_ub_n),
        .psram_zz_n(psram_zz_n)
    );

    psram_model #(
        .ADDR_WIDTH(ADDR_WIDTH),
        .DATA_WIDTH(DATA_WIDTH),
        .DEPTH(4096)
    ) memory (
        .clk   (clk),
        .a     (psram_a),
        .dq    (psram_dq),
        .ce_n  (psram_ce_n),
        .oe_n  (psram_oe_n),
        .we_n  (psram_we_n),
        .lb_n  (psram_lb_n),
        .ub_n  (psram_ub_n),
        .zz_n  (psram_zz_n)
    );

    // ============================================================
    // Clock
    // ============================================================

    initial begin
        clk = 1'b0;
        forever #(CLK_PERIOD / 2.0) clk = ~clk;
    end

    initial begin
        $dumpfile("sim/psram_page_mode.vcd");
        $dumpvars(0, tb);
    end

    // ============================================================
    // CE# pulse counter -- counts how many times psram_ce_n rises
    // (i.e. how many times a session actually closed), so tests
    // can check that a page really stayed open (or really closed)
    // without hand-parsing the DUT's internal state.
    // ============================================================

    integer ce_close_count;
    integer ce_close_before_burst;

    always @(posedge psram_ce_n)
        ce_close_count = ce_close_count + 1;

    // ============================================================
    // Helpers
    // ============================================================

    real req_time;
    real latency_ns;

    task write_word;
        input [ADDR_WIDTH-1:0] addr;
        input [DATA_WIDTH-1:0] data;
        begin
            @(posedge clk);
            mem_addr  <= addr;
            mem_wdata <= data;
            mem_wr    <= 1'b1;
            mem_lb_n  <= 1'b0;
            mem_ub_n  <= 1'b0;
            mem_req   <= 1'b1;

            @(posedge clk);
            mem_req <= 1'b0;

            wait (mem_ready);
            @(posedge clk);
        end
    endtask

    // Full-word read, records latency (mem_req assertion -> mem_ready)
    // in latency_ns for the caller to inspect.
    task read_word;
        input  [ADDR_WIDTH-1:0] addr;
        input  [DATA_WIDTH-1:0] expected;
        begin
            @(posedge clk);
            mem_addr <= addr;
            mem_wr   <= 1'b0;
            mem_lb_n <= 1'b0;
            mem_ub_n <= 1'b0;
            mem_req  <= 1'b1;

            req_time = $realtime;

            @(posedge clk);
            mem_req <= 1'b0;

            wait (mem_ready);
            latency_ns = $realtime - req_time;

            if (mem_rdata !== expected) begin
                $display("READ addr=0x%06x FAIL got=0x%04x expected=0x%04x",
                          addr, mem_rdata, expected);
                $fatal;
            end

            $display("READ addr=0x%06x data=0x%04x latency=%0.1fns PASS",
                      addr, mem_rdata, latency_ns);

            @(posedge clk);
        end
    endtask

    task read_word_be;
        input  [ADDR_WIDTH-1:0] addr;
        input                   lb;
        input                   ub;
        input  [DATA_WIDTH-1:0] expected;
        begin
            @(posedge clk);
            mem_addr <= addr;
            mem_wr   <= 1'b0;
            mem_lb_n <= lb;
            mem_ub_n <= ub;
            mem_req  <= 1'b1;

            req_time = $realtime;

            @(posedge clk);
            mem_req <= 1'b0;

            wait (mem_ready);
            latency_ns = $realtime - req_time;

            if (mem_rdata !== expected) begin
                $display("READ(BE) addr=0x%06x FAIL got=0x%04x expected=0x%04x",
                          addr, mem_rdata, expected);
                $fatal;
            end

            $display("READ(BE) addr=0x%06x LB#=%b UB#=%b data=0x%04x latency=%0.1fns PASS",
                      addr, lb, ub, mem_rdata, latency_ns);

            @(posedge clk);
        end
    endtask

    task expect_ce_closes;
        input integer expected_count;
        input [8*48-1:0] label;
        begin
            if (ce_close_count !== expected_count) begin
                $display("CE# close-count FAIL (%0s): got=%0d expected=%0d",
                          label, ce_close_count, expected_count);
                $fatal;
            end else begin
                $display("CE# close-count OK (%0s): %0d", label, ce_close_count);
            end
        end
    endtask

    // ============================================================
    // Test
    // ============================================================

    integer i;
    reg [DATA_WIDTH-1:0] page_data [0:15];

    initial begin

        mem_req   = 1'b0;
        mem_wr    = 1'b0;
        mem_addr  = 0;
        mem_wdata = 0;
        mem_lb_n  = 1'b1;
        mem_ub_n  = 1'b1;

        ce_close_count = 0;

        rst = 1'b1;
        repeat (5) @(posedge clk);
        rst = 1'b0;

        $display("");
        $display("========================================");
        $display("PSRAM PAGE MODE TEST");
        $display("%0d MHz -- ACCESS_CYCLES=%0d PAGE_CYCLES=%0d",
                  CLK_FREQ_MHZ, ACCESS_CYCLES, PAGE_CYCLES);
        $display("========================================");
        $display("");

        // Wait through STATE_INIT + the CR software-access-sequence
        // (2 dummy reads + 2 writes at the top address) -- if the
        // sequence violates any read/write timing the strict
        // psram_model will $fatal before we ever get here.

        wait (dut.state == dut.STATE_IDLE);
        $display("PSRAM init + CR page-mode enable sequence complete");
        $display("");

        // Reset the close-counter here: boot (reset release + the
        // 4-step CR sequence) legitimately toggles CE# several
        // times and that's not what the test below is checking.
        ce_close_count = 0;

        // ========================================================
        // Fill one 16-word page (addresses share bits above A[3])
        // plus one word in the next page, for boundary testing.
        // ========================================================

        for (i = 0; i < 16; i = i + 1) begin
            page_data[i] = 16'hA000 + i[15:0];
            write_word(23'h000100 + i, page_data[i]);
        end

        write_word(23'h000110, 16'hB000); // first word of the NEXT page

        expect_ce_closes(17, "after 17 writes");

        // ========================================================
        // Same-page sequential reads: first word pays full tAA,
        // every following word in the same page must be a fast
        // PAGE_CYCLES continuation with CE# never toggling.
        // ========================================================

        $display("");
        $display("---- same-page sequential read burst ----");

        read_word(23'h000100, page_data[0]);

        if (latency_ns < ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: first word of a fresh page was faster than tAA (%0.1fns < %0.1fns)",
                      latency_ns, ACCESS_CYCLES * CLK_PERIOD);
            $fatal;
        end

        // Opening a page (a completed read) does NOT close CE# --
        // that's the whole point, the session stays open.
        expect_ce_closes(17, "page opened, CE# still low");

        for (i = 1; i < 16; i = i + 1) begin
            read_word(23'h000100 + i, page_data[i]);

            if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
                $display("FAIL: same-page word %0d did not use the fast path (%0.1fns >= tAA %0.1fns)",
                          i, latency_ns, ACCESS_CYCLES * CLK_PERIOD);
                $fatal;
            end

            if (latency_ns > PAGE_CYCLES * CLK_PERIOD + CLK_PERIOD) begin
                $display("FAIL: same-page word %0d slower than expected (%0.1fns)",
                          i, latency_ns);
                $fatal;
            end
        end

        // CE# must NOT have toggled again across the whole 16-word
        // burst -- it should still be exactly one open session.
        expect_ce_closes(17, "still one open session after 16-word burst");

        $display("PAGE MODE BURST SPEEDUP CONFIRMED");
        $display("");

        // ========================================================
        // Page-boundary crossing while CE# stays open: a READ to a
        // different page still doesn't need to close CE# (only a
        // WRITE does) -- but that one word pays the full tAA per
        // the datasheet rule ("any change in addresses A[4] or
        // higher initiates a new tAA access time"), then
        // continuations in the new page are fast again.
        // ========================================================

        $display("---- page-boundary crossing (still open) ----");

        read_word(23'h000110, 16'hB000);

        if (latency_ns < ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: page-crossing word was not full tAA (%0.1fns)", latency_ns);
            $fatal;
        end

        expect_ce_closes(17, "page crossing did not toggle CE#");

        // A WRITE, unlike a READ, always closes the page -- and it
        // costs two CE# pulses: one to close the read session that
        // was open, one for the write's own transaction.
        write_word(23'h000111, 16'hB001);
        expect_ce_closes(19, "write after page crossing closes the session");

        read_word(23'h000111, 16'hB001);
        expect_ce_closes(19, "fresh read reopens its own session (no close)");

        read_word(23'h000112, 16'h0000); // untouched location -> reset value
        if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: second word of new page-crossing session should be fast");
            $fatal;
        end
        expect_ce_closes(19, "same-page continuation after the crossing");

        $display("");

        // ========================================================
        // A WRITE mid-session must close the page.
        // ========================================================

        $display("---- write closes an open page ----");

        // Still the SAME open session as above (page 0x11) -- a
        // READ to a different page (0x10) is just another
        // page-miss continuation, not a close.
        read_word(23'h000100, page_data[0]);
        expect_ce_closes(19, "page-miss read continuation (still open)");

        read_word(23'h000101, page_data[1]);   // fast continuation
        if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: continuation before the write should be fast");
            $fatal;
        end
        expect_ce_closes(19, "still open before the write");

        write_word(23'h000200, 16'hC0DE);
        expect_ce_closes(21, "write forces a close (2 pulses: close + write)");

        read_word(23'h000200, 16'hC0DE);
        expect_ce_closes(21, "post-write read opens a fresh session (no close)");

        $display("");

        // ========================================================
        // Byte-enable changes must NOT close the page.
        //
        // int8_memory_access.v alternates LB#/UB# on essentially
        // every access (byte-granular reads over the 16-bit PSRAM
        // bus, addr[0] selects the byte) -- this is the actual
        // real-world access pattern (e.g. graph_engine's edge-list
        // gather), so it must stay on the fast page-hit path, not
        // force a close on every single byte.
        // ========================================================

        $display("---- byte-enable changes stay on the fast path ----");

        read_word(23'h000102, page_data[2]);
        expect_ce_closes(21, "page-miss read continuation, still open");

        read_word_be(23'h000102, 1'b0, 1'b1, {8'h00, page_data[2][7:0]});
        expect_ce_closes(21, "LB# only -- still open, still fast");
        if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: byte-enable-only change should stay on the fast path");
            $fatal;
        end

        read_word_be(23'h000102, 1'b1, 1'b0, {page_data[2][15:8], 8'h00});
        expect_ce_closes(21, "UB# only -- still open, still fast");
        if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: byte-enable-only change should stay on the fast path");
            $fatal;
        end

        read_word_be(23'h000103, 1'b0, 1'b1, {8'h00, page_data[3][7:0]});
        expect_ce_closes(21, "new address + byte-enable change together -- still fast");
        if (latency_ns >= ACCESS_CYCLES * CLK_PERIOD) begin
            $display("FAIL: address+byte-enable change together should stay fast");
            $fatal;
        end

        $display("");

        // ========================================================
        // tCEM safety timeout: an open page with no further
        // requests must close itself well before the 8us CE#-low
        // refresh limit (PAGE_TIMEOUT_CYCLES, ~6us of margin) --
        // not just "eventually", but on its own, unprompted.
        // ========================================================

        $display("---- tCEM idle timeout closes an unattended open page ----");

        read_word(23'h000104, page_data[4]); // still page 0x10, still open
        expect_ce_closes(21, "same page continuation, waiting idle now");

        repeat (dut.PAGE_TIMEOUT_CYCLES + 4) @(posedge clk);

        expect_ce_closes(22, "idle page auto-closed by the tCEM timeout");

        read_word(23'h000104, page_data[4]); // must still work correctly
        expect_ce_closes(22, "fresh session opened, not yet closed again");

        $display("");

        // ========================================================
        // tCEM budget mid-burst: a long run of back-to-back
        // same-page HITS (never idle, never a write) must still be
        // split before the limit -- not just the idle-timeout case
        // above. sim/psram_model.v independently enforces the real
        // 8us tCEM hard limit ($fatal on violation); this is a real
        // safety net, not a rubber stamp, so a passing run here is
        // genuine proof the RTL splits the burst in time, with
        // margin, not just "in simulation it happened to work".
        // ========================================================

        $display("---- tCEM budget forces a split mid-burst (never idle) ----");

        for (i = 0; i < 16; i = i + 1) begin
            page_data[i] = 16'hC000 + i[15:0];
            write_word(23'h000500 + i, page_data[i]);
        end

        read_word(23'h000500, page_data[0]); // opens page 0x50
        ce_close_before_burst = ce_close_count;

        // 2*PAGE_TIMEOUT_CYCLES/16 round trips through this 16-word
        // page comfortably crosses PAGE_TIMEOUT_CYCLES worth of
        // STATE_READ time (each hit costs PAGE_CYCLES inside
        // STATE_READ, plus this task's own idle cycles between
        // requests, both counted by hold_cycles) while never once
        // idling long enough on its own to hit the separate
        // idle-timeout path above -- this is the "busy" case.
        for (i = 0; i < 2 * dut.PAGE_TIMEOUT_CYCLES; i = i + 1)
            read_word(23'h000500 + (i % 16), page_data[i % 16]);

        if (ce_close_count <= ce_close_before_burst) begin
            $display("FAIL: expected at least one mid-burst split, CE# never toggled (before=%0d after=%0d)",
                      ce_close_before_burst, ce_close_count);
            $fatal;
        end

        $display("CE# split during long burst confirmed: %0d -> %0d closes",
                  ce_close_before_burst, ce_close_count);
        $display("(no tCEM $fatal from psram_model.v -- split happened with margin)");
        $display("");

        // ========================================================
        // Data-integrity stress: random-ish scattered pages,
        // mixing writes and page-local read bursts.
        // ========================================================

        $display("---- mixed stress: scattered pages + local bursts ----");

        for (i = 0; i < 64; i = i + 1) begin
            reg [ADDR_WIDTH-1:0] base;
            reg [DATA_WIDTH-1:0] val;
            integer j;

            base = ((i * 8191) ^ (i << 6)) & 23'h000FF0; // page-aligned
            val  = (i * 733) ^ 16'h5A5A;

            for (j = 0; j < 4; j = j + 1)
                write_word(base + j, val + j[15:0]);

            for (j = 0; j < 4; j = j + 1)
                read_word(base + j, val + j[15:0]);
        end

        $display("");
        $display("========================================");
        $display("PSRAM PAGE MODE TEST PASSED");
        $display("========================================");
        $display("");

        $finish;
    end

endmodule