`timescale 1ns/1ps
// Self-checking simulation, separate from the synthesizable two RTL modules.
module tb_systolic_2x2;
  logic clk=0;
  always #5 clk=~clk;
  logic reset=0, clear=0, advance=1;
  logic signed [7:0] a_left0=0,a_left1=0,b_top0=0,b_top1=0;
  logic [1:0] a_valid=0,b_valid=0;
  logic a0_k=0,a1_k=0,b0_k=0,b1_k=0;
  wire signed [16:0] c00,c01,c10,c11;
  wire done,protocol_error;
  logic signed [7:0] A[0:1][0:1],B[0:1][0:1];
  integer checks=0;
  systolic_2x2 dut(.*);

  // A K=4 PE also checks the accumulator boundary 65536 in AW=18.
  logic p_reset=0,p_clear=0,p_advance=0,p_av=0,p_bv=0;
  logic signed [7:0] p_a=0,p_b=0;
  logic [1:0] p_ak=0,p_bk=0;
  wire signed [17:0] p_acc;
  wire [2:0] p_count;
  wire p_done,p_error;
  signed_mac_pe #(.W(8),.K(4),.AW(18)) boundary_pe(
    .clk(clk),.reset(p_reset),.clear(p_clear),.advance(p_advance),
    .a_in(p_a),.b_in(p_b),.a_valid_in(p_av),.b_valid_in(p_bv),
    .a_k_in(p_ak),.b_k_in(p_bk),.a_out(),.b_out(),
    .a_valid_out(),.b_valid_out(),.a_k_out(),.b_k_out(),
    .acc(p_acc),.count(p_count),.done(p_done),.protocol_error(p_error)
  );

  task automatic expect_matrix(input integer x00,x01,x10,x11,input bit finished);
    begin
      checks=checks+1;
      if ($signed(c00)!==x00 || $signed(c01)!==x01 ||
          $signed(c10)!==x10 || $signed(c11)!==x11 || done!==finished || protocol_error!==0)
        $fatal(1,"matrix mismatch got [%0d,%0d;%0d,%0d] done=%b err=%b expected [%0d,%0d;%0d,%0d] done=%b",
          c00,c01,c10,c11,done,protocol_error,x00,x01,x10,x11,finished);
    end
  endtask

  task automatic clear_grid;
    begin
      @(negedge clk); clear=1; advance=1; a_valid=0; b_valid=0;
      @(posedge clk); #1; expect_matrix(0,0,0,0,0);
      if(dut.count00!==0||dut.count01!==0||dut.count10!==0||dut.count11!==0)
        $fatal(1,"clear did not zero counts");
      @(negedge clk); clear=0;
    end
  endtask

  task automatic step(input integer t);
    begin
      // No extra clock is consumed between clear_grid and this task.
      a_valid=0;b_valid=0;a_left0=0;a_left1=0;b_top0=0;b_top1=0;
      a0_k=0;a1_k=0;b0_k=0;b1_k=0;advance=1;
      case(t)
        0: begin a_left0=A[0][0];b_top0=B[0][0];a_valid=2'b01;b_valid=2'b01;end
        1: begin a_left0=A[0][1];a_left1=A[1][0];b_top0=B[1][0];b_top1=B[0][1];
                 a_valid=2'b11;b_valid=2'b11;a0_k=1;b0_k=1;end
        2: begin a_left1=A[1][1];b_top1=B[1][1];a_valid=2'b10;b_valid=2'b10;a1_k=1;b1_k=1;end
        3: begin end
        default: $fatal(1,"invalid step");
      endcase
      @(posedge clk); #1;
      @(negedge clk);
    end
  endtask

  function automatic [255:0] state_signature;
    begin
      state_signature={c00,c01,c10,c11,
        dut.pe00.a_out,dut.pe00.b_out,dut.pe01.a_out,dut.pe01.b_out,
        dut.pe10.a_out,dut.pe10.b_out,dut.pe11.a_out,dut.pe11.b_out,
        dut.pe00.a_valid_out,dut.pe00.b_valid_out,dut.pe01.a_valid_out,dut.pe01.b_valid_out,
        dut.pe10.a_valid_out,dut.pe10.b_valid_out,dut.pe11.a_valid_out,dut.pe11.b_valid_out,
        dut.pe00.a_k_out,dut.pe00.b_k_out,dut.pe01.a_k_out,dut.pe01.b_k_out,
        dut.pe10.a_k_out,dut.pe10.b_k_out,dut.pe11.a_k_out,dut.pe11.b_k_out,
        dut.count00,dut.count01,dut.count10,dut.count11,done,protocol_error};
    end
  endfunction

  task automatic pause_twice;
    reg [255:0] held;
    integer q;
    begin
      held=state_signature();advance=0;
      for(q=0;q<2;q=q+1)begin
        @(posedge clk);#1;checks=checks+1;
        if(state_signature()!==held)$fatal(1,"state changed while advance=0");
        @(negedge clk);
      end
      advance=1;
    end
  endtask

  task automatic boundary_case(input integer avalue,bvalue,expected_sum);
    integer k;
    begin
      @(negedge clk);p_clear=1;p_advance=0;p_av=0;p_bv=0;
      @(posedge clk);#1;if(p_acc!==0||p_count!==0||p_error!==0)$fatal(1,"PE clear failed");
      @(negedge clk);p_clear=0;p_advance=1;p_av=1;p_bv=1;p_a=avalue;p_b=bvalue;
      for(k=0;k<4;k=k+1)begin
        p_ak=k;p_bk=k;@(posedge clk);#1;
        checks=checks+1;
        if($signed(p_acc)!==(k+1)*avalue*bvalue||p_count!==(k+1)||p_error!==0)
          $fatal(1,"signed boundary prefix failed k=%0d acc=%0d",k,p_acc);
        @(negedge clk);
      end
      if($signed(p_acc)!==expected_sum||p_done!==1)$fatal(1,"signed boundary final failed");
      p_advance=0;p_av=0;p_bv=0;
    end
  endtask

  initial begin
    // Synchronous reset is checked while advance=0.
    advance=0;reset=1;p_reset=1;
    @(posedge clk);#1;expect_matrix(0,0,0,0,0);
    @(negedge clk);reset=0;p_reset=0;
    A[0][0]=1;A[0][1]=2;A[1][0]=3;A[1][1]=4;
    B[0][0]=5;B[0][1]=6;B[1][0]=7;B[1][1]=8;
    clear_grid();
    step(0);expect_matrix(5,0,0,0,0);
    step(1);expect_matrix(19,6,15,0,0);
    step(2);expect_matrix(19,22,43,18,0);
    step(3);expect_matrix(19,22,43,50,1);
    if(dut.count00!==2||dut.count01!==2||dut.count10!==2||dut.count11!==2)
      $fatal(1,"not all PEs consumed K=2");
    $display("TRACE E0..E3: [5,0;0,0], [19,6;15,0], [19,22;43,18], [19,22;43,50]");

    clear_grid();step(0);pause_twice();expect_matrix(5,0,0,0,0);
    step(1);step(2);step(3);expect_matrix(19,22,43,50,1);
    $display("PAUSE: four advance edges plus two paused edges = six compute edges");

    A[0][0]=2;A[0][1]=-1;A[1][0]=0;A[1][1]=3;
    B[0][0]=4;B[0][1]=5;B[1][0]=-2;B[1][1]=1;
    clear_grid();step(0);step(1);step(2);step(3);expect_matrix(10,9,-6,3,1);

    // Reset aborts a partially completed array and clears forwarded validity.
    clear_grid();step(0);advance=0;reset=1;
    @(posedge clk);#1;expect_matrix(0,0,0,0,0);
    if(dut.av01!==0||dut.av11!==0||dut.bv10!==0||dut.bv11!==0)
      $fatal(1,"reset left valid tokens");
    @(negedge clk);reset=0;

    boundary_case(-128,-128,65536);
    boundary_case(-128,127,-65024);
    boundary_case(127,127,64516);
    // A mismatched k pair must raise an error without consuming a MAC.
    @(negedge clk);p_clear=1;
    @(posedge clk);#1;
    @(negedge clk);p_clear=0;p_advance=1;p_av=1;p_bv=1;p_a=2;p_b=3;p_ak=0;p_bk=1;
    @(posedge clk);#1;checks=checks+1;
    if(p_error!==1||p_count!==0||p_acc!==0)$fatal(1,"tag mismatch was accepted");
    $display("RTL PASS: %0d self-checks; signed width, skew, count, pause, reset, tag rejection",checks);
    $finish;
  end

  initial begin #10000;$fatal(1,"watchdog timeout");end
endmodule