`timescale 1ns/1ps

// ============================================================
// EXP-0089 -- isolated correctness test for systolic_group.v: one
// group of 4 packed_pe.v instances sharing ONE real weight fetch
// (broadcast, not a literal PE-to-PE systolic shift register -- see
// systolic_group.v's own header for the real, user-confirmed design
// choice). Same real discipline as every other new module in this
// project (act_tile_fetch.v EXP-0079, ddr_prefetch_mgr.v EXP-0083,
// result_writeback.v EXP-0088): verify in isolation, with a real
// memory backend and a real, independently-reproduced golden model,
// BEFORE any Director/SPI-protocol integration.
//
// Real backend: burst_mem_model32.v (same EXP-0084 model every other
// v3 isolated testbench uses) + sdram_arbiter_n.v with NUM_REQ=5 (1
// group-level weight-fetch requester + 4 independent per-PE
// activation-fetch/writeback requesters) -- sdram_arbiter_n.v's own
// NUM_REQ already generalizes to this without any change, confirmed
// by direct reuse here, not by inspection.
//
// Runs TWO consecutive group jobs (different positions/weights each
// time) specifically to catch any "forgot to clear a per-job latch"
// bug in the group's own barrier state (pe_acked/pe_done_latch) --
// a single-job test would not exercise that reset path at all.
// ============================================================
module tb;
    localparam BURST_LEN  = 8;
    localparam SDRAM_ADDR_WIDTH = 25;
    localparam CLK_FREQ_MHZ = 64;
    localparam CLK_PERIOD_NS = 1000.0/CLK_FREQ_MHZ;

    localparam DATA_WIDTH  = 8;
    localparam P_IN        = 8;
    localparam ACC_WIDTH   = 32;
    localparam ADDR_WIDTH  = 26;
    localparam N_INPUTS    = 128;
    localparam N_TILES     = N_INPUTS/P_IN;
    localparam LAYER_BYTES = N_INPUTS;
    localparam WORDS_PER_LAYER = LAYER_BYTES/2;
    localparam NUM_REQ = 5; // 1 group weight-fetch + 4 PE activation/writeback

    reg clk = 0;
    always #(CLK_PERIOD_NS/2.0) clk = ~clk;
    reg rst;
    integer cyc;
    always @(posedge clk) if (!rst) cyc <= cyc + 1;

    // ---- real burst-memory backend, shared via a real 5-way arbiter ----
    wire                          ctrl_req, ctrl_wr;
    wire [SDRAM_ADDR_WIDTH-1:0]   ctrl_addr;
    wire [32*BURST_LEN-1:0]       ctrl_wdata;
    wire [4*BURST_LEN-1:0]        ctrl_wmask;
    wire [32*BURST_LEN-1:0]       ctrl_rdata;
    wire ctrl_ready, ctrl_busy;

    reg  wpre_req, wpre_wr;
    reg  [SDRAM_ADDR_WIDTH-1:0] wpre_addr;
    reg  [32*BURST_LEN-1:0] wpre_wdata;
    reg  pre_active;

    wire [NUM_REQ-1:0]              arb_active, arb_grant, arb_req, arb_wr;
    wire [NUM_REQ*SDRAM_ADDR_WIDTH-1:0] arb_addr;
    wire [NUM_REQ*32*BURST_LEN-1:0] arb_wdata, arb_rdata;
    wire [NUM_REQ*4*BURST_LEN-1:0]  arb_wmask;
    wire [NUM_REQ-1:0]              arb_ready, arb_busy;

    wire real_ctrl_req, real_ctrl_wr;
    wire [SDRAM_ADDR_WIDTH-1:0] real_ctrl_addr;
    wire [32*BURST_LEN-1:0] real_ctrl_wdata;
    wire [4*BURST_LEN-1:0] real_ctrl_wmask;

    sdram_arbiter_n #(.NUM_REQ(NUM_REQ), .ADDR_WIDTH(SDRAM_ADDR_WIDTH), .BURST_LEN(BURST_LEN)) u_arb (
        .clk(clk), .rst(rst),
        .req_active(arb_active), .req_grant(arb_grant),
        .req_req(arb_req), .req_wr(arb_wr), .req_addr(arb_addr),
        .req_wdata(arb_wdata), .req_wmask(arb_wmask),
        .req_rdata(arb_rdata), .req_ready(arb_ready), .req_busy(arb_busy),
        .ctrl_req(real_ctrl_req), .ctrl_wr(real_ctrl_wr), .ctrl_addr(real_ctrl_addr),
        .ctrl_wdata(real_ctrl_wdata), .ctrl_wmask(real_ctrl_wmask),
        .ctrl_rdata(ctrl_rdata), .ctrl_ready(ctrl_ready), .ctrl_busy(ctrl_busy)
    );

    assign ctrl_req   = pre_active ? wpre_req   : real_ctrl_req;
    assign ctrl_wr    = pre_active ? wpre_wr    : real_ctrl_wr;
    assign ctrl_addr  = pre_active ? wpre_addr  : real_ctrl_addr;
    assign ctrl_wdata = pre_active ? wpre_wdata : real_ctrl_wdata;
    assign ctrl_wmask = pre_active ? {(4*BURST_LEN){1'b0}} : real_ctrl_wmask;

    burst_mem_model32 #(
        .BURST_LEN(BURST_LEN), .ADDR_WIDTH(SDRAM_ADDR_WIDTH)
    ) u_mem (
        .clk(clk), .rst(rst),
        .req(ctrl_req), .wr(ctrl_wr), .addr(ctrl_addr), .wdata(ctrl_wdata), .wmask(ctrl_wmask),
        .rdata(ctrl_rdata), .ready(ctrl_ready), .busy(ctrl_busy)
    );

    function automatic signed [7:0] weight_byte(input integer li, input integer t);
        reg [7:0] tmp;
        begin
            tmp = (li*17 + t*29 + 13) & 8'hFF;
            weight_byte = $signed(tmp);
        end
    endfunction
    function automatic signed [7:0] input_byte(input integer li, input integer pos, input integer t);
        reg [7:0] tmp;
        begin
            tmp = (li*11 + pos*41 + t*7 + 3) & 8'hFF;
            input_byte = $signed(tmp);
        end
    endfunction

    task automatic sdram_write_burst(input [SDRAM_ADDR_WIDTH-1:0] word_addr, input [32*BURST_LEN-1:0] data);
        begin
            @(posedge clk); while (ctrl_busy) @(posedge clk);
            wpre_req = 1'b1; wpre_wr = 1'b1; wpre_addr = word_addr; wpre_wdata = data;
            @(posedge clk); wpre_req = 1'b0;
            while (!ctrl_ready) @(posedge clk);
        end
    endtask

    task automatic preload_sdram_layer(input integer li);
        integer bi, wb, tt;
        reg [32*BURST_LEN-1:0] burst_data;
        begin
            for (bi = 0; bi < (LAYER_BYTES/(4*BURST_LEN)); bi = bi + 1) begin
                for (wb = 0; wb < BURST_LEN; wb = wb + 1) begin
                    tt = bi*(4*BURST_LEN) + wb*4;
                    burst_data[wb*32 +: 32] = {weight_byte(li, tt+3), weight_byte(li, tt+2),
                                                weight_byte(li, tt+1), weight_byte(li, tt)};
                end
                sdram_write_burst((li*WORDS_PER_LAYER + bi*BURST_LEN), burst_data);
            end
        end
    endtask

    localparam [ADDR_WIDTH-1:0] ACT_MEM_BASE = 26'h10000;
    function automatic [ADDR_WIDTH-1:0] act_x_base(input integer li, input integer pos);
        act_x_base = ACT_MEM_BASE + (li*16 + pos) * ((N_TILES/4)*BURST_LEN);
    endfunction

    task automatic preload_sdram_activation(input integer li, input integer pos);
        integer tq, qi;
        reg [32*BURST_LEN-1:0] burst_data;
        reg [ADDR_WIDTH-1:0] base;
        begin
            base = act_x_base(li, pos);
            for (tq = 0; tq < N_TILES/4; tq = tq + 1) begin
                burst_data = {(32*BURST_LEN){1'b0}};
                for (qi = 0; qi < 4; qi = qi + 1)
                    burst_data[qi*64 +: 64] = {input_byte(li, pos, (4*tq+qi)*P_IN + 7), input_byte(li, pos, (4*tq+qi)*P_IN + 6),
                                                input_byte(li, pos, (4*tq+qi)*P_IN + 5), input_byte(li, pos, (4*tq+qi)*P_IN + 4),
                                                input_byte(li, pos, (4*tq+qi)*P_IN + 3), input_byte(li, pos, (4*tq+qi)*P_IN + 2),
                                                input_byte(li, pos, (4*tq+qi)*P_IN + 1), input_byte(li, pos, (4*tq+qi)*P_IN + 0)};
                sdram_write_burst(base[SDRAM_ADDR_WIDTH-1:0] + tq*BURST_LEN, burst_data);
            end
        end
    endtask

    // ---- systolic_group.v (DUT) ----
    reg  job_start;
    reg  [ADDR_WIDTH-1:0] w_base;
    reg  [15:0] n_tiles_in;
    reg  [4*ADDR_WIDTH-1:0] pe_x_base_a, pe_x_base_b, pe_result_addr_a, pe_result_addr_b;
    reg  [4*16-1:0] pe_node_id_a, pe_node_id_b;
    wire job_done;
    wire [4*DATA_WIDTH-1:0] pe_result_data_a, pe_result_data_b;
    wire [4*16-1:0] pe_result_node_id_a, pe_result_node_id_b;
    wire [4*ADDR_WIDTH-1:0] pe_result_addr_a_out, pe_result_addr_b_out;

    wire                    grp_mem_active;
    wire [3:0]              pe_mem_active;

    systolic_group #(
        .DATA_WIDTH(DATA_WIDTH), .P_IN(P_IN), .ACC_WIDTH(ACC_WIDTH),
        .BURST_LEN(BURST_LEN), .ADDR_WIDTH(ADDR_WIDTH), .LAYER_BYTES(LAYER_BYTES)
    ) dut (
        .clk(clk), .rst(rst),
        .job_start(job_start), .w_base(w_base), .n_tiles(n_tiles_in),
        .pe_x_base_a(pe_x_base_a), .pe_x_base_b(pe_x_base_b),
        .pe_result_addr_a(pe_result_addr_a), .pe_result_addr_b(pe_result_addr_b),
        .pe_node_id_a(pe_node_id_a), .pe_node_id_b(pe_node_id_b),
        .job_done(job_done),
        .pe_result_data_a(pe_result_data_a), .pe_result_data_b(pe_result_data_b),
        .pe_result_node_id_a(pe_result_node_id_a), .pe_result_node_id_b(pe_result_node_id_b),
        .pe_result_addr_a_out(pe_result_addr_a_out), .pe_result_addr_b_out(pe_result_addr_b_out),
        .mem_active(grp_mem_active), .mem_grant(arb_grant[0]),
        .ctrl_req(arb_req[0]), .ctrl_wr(arb_wr[0]), .ctrl_addr(arb_addr[0*SDRAM_ADDR_WIDTH +: SDRAM_ADDR_WIDTH]),
        .ctrl_wdata(arb_wdata[0*32*BURST_LEN +: 32*BURST_LEN]), .ctrl_wmask(arb_wmask[0*4*BURST_LEN +: 4*BURST_LEN]),
        .ctrl_rdata(arb_rdata[0*32*BURST_LEN +: 32*BURST_LEN]), .ctrl_ready(arb_ready[0]), .ctrl_busy(arb_busy[0]),
        .pe_mem_active(pe_mem_active), .pe_mem_grant(arb_grant[4:1]),
        .pe_ctrl_req(arb_req[4:1]), .pe_ctrl_wr(arb_wr[4:1]),
        .pe_ctrl_addr(arb_addr[1*SDRAM_ADDR_WIDTH +: 4*SDRAM_ADDR_WIDTH]),
        .pe_ctrl_wdata(arb_wdata[1*32*BURST_LEN +: 4*32*BURST_LEN]),
        .pe_ctrl_wmask(arb_wmask[1*4*BURST_LEN +: 4*4*BURST_LEN]),
        .pe_ctrl_rdata(arb_rdata[1*32*BURST_LEN +: 4*32*BURST_LEN]),
        .pe_ctrl_ready(arb_ready[4:1]), .pe_ctrl_busy(arb_busy[4:1])
    );

    assign arb_active[0]   = grp_mem_active;
    assign arb_active[4:1] = pe_mem_active;

    integer errors, tests;
    integer li_i, gi, wd;
    integer acc, s;
    reg signed [DATA_WIDTH-1:0] expected [0:7]; // 4 PEs x 2 lanes

    task automatic run_group_job(input integer li, input integer pos_base);
        integer pe, lane, pos, tt;
        reg [DATA_WIDTH-1:0] got_a, got_b;
        begin
            tests = tests + 1;
            @(posedge clk);
            job_start  = 1'b1;
            w_base     = li*WORDS_PER_LAYER;
            n_tiles_in = N_TILES[15:0];
            for (pe = 0; pe < 4; pe = pe + 1) begin
                pe_x_base_a[pe*ADDR_WIDTH +: ADDR_WIDTH] = act_x_base(li, pos_base + pe*2);
                pe_x_base_b[pe*ADDR_WIDTH +: ADDR_WIDTH] = act_x_base(li, pos_base + pe*2 + 1);
                pe_result_addr_a[pe*ADDR_WIDTH +: ADDR_WIDTH] = 26'h9000 + pe*2;
                pe_result_addr_b[pe*ADDR_WIDTH +: ADDR_WIDTH] = 26'h9000 + pe*2 + 1;
                pe_node_id_a[pe*16 +: 16] = li*100 + pos_base + pe*2;
                pe_node_id_b[pe*16 +: 16] = li*100 + pos_base + pe*2 + 1;
            end
            @(posedge clk);
            job_start = 1'b0;

            // real, independently-reproduced golden model (same formula
            // as tb_packed_slot.v's own, applied per PE/lane)
            for (pe = 0; pe < 4; pe = pe + 1) begin
                for (lane = 0; lane < 2; lane = lane + 1) begin
                    pos = pos_base + pe*2 + lane;
                    acc = 0;
                    for (tt = 0; tt < N_INPUTS; tt = tt + 1)
                        acc = acc + (input_byte(li, pos, tt) * weight_byte(li, tt));
                    s = acc;
                    if (s <= 0) expected[pe*2+lane] = 0;
                    else if (s > 127) expected[pe*2+lane] = 8'sd127;
                    else expected[pe*2+lane] = s[DATA_WIDTH-1:0];
                end
            end

            wd = 0;
            while (!job_done && wd < 4000) begin @(posedge clk); wd = wd + 1; end
            if (!job_done) begin
                $display("FAIL li=%0d pos_base=%0d: TIMEOUT waiting for group job_done", li, pos_base);
                errors = errors + 1;
            end else begin
                for (pe = 0; pe < 4; pe = pe + 1) begin
                    got_a = pe_result_data_a[pe*DATA_WIDTH +: DATA_WIDTH];
                    got_b = pe_result_data_b[pe*DATA_WIDTH +: DATA_WIDTH];
                    if (got_a !== expected[pe*2] || got_b !== expected[pe*2+1]) begin
                        $display("FAIL li=%0d pos_base=%0d PE%0d: got_a=%0d got_b=%0d expected_a=%0d expected_b=%0d",
                                  li, pos_base, pe, $signed(got_a), $signed(got_b),
                                  $signed(expected[pe*2]), $signed(expected[pe*2+1]));
                        errors = errors + 1;
                    end else begin
                        $display("PASS li=%0d pos_base=%0d PE%0d: a=%0d b=%0d (systolic_group.v)",
                                  li, pos_base, pe, $signed(got_a), $signed(got_b));
                    end
                end
            end
        end
    endtask

    initial begin
        errors = 0; tests = 0; cyc = 0;
        rst = 1; pre_active = 1'b1;
        wpre_req = 0; wpre_wr = 0; wpre_addr = 0; wpre_wdata = 0;
        job_start = 0; w_base = 0; n_tiles_in = 0;
        pe_x_base_a = 0; pe_x_base_b = 0; pe_result_addr_a = 0; pe_result_addr_b = 0;
        pe_node_id_a = 0; pe_node_id_b = 0;
        repeat(5) @(posedge clk);
        rst = 0;
        @(posedge clk); while (ctrl_busy) @(posedge clk);

        $display("=== preload SDRAM: 2 layers' weights + 16 activation positions ===");
        for (li_i = 0; li_i < 2; li_i = li_i + 1) begin
            preload_sdram_layer(li_i);
            for (gi = 0; gi < 8; gi = gi + 1) preload_sdram_activation(li_i, gi);
        end
        @(posedge clk);
        pre_active = 1'b0;

        $display("=== systolic_group.v: 2 group jobs (4 PEs x 2 lanes each) ===");
        run_group_job(0, 0);
        run_group_job(1, 0);

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