`timescale 1ns/1ps
module tb_cnn_stream;
  logic clk = 0;
  always #5 clk = !clk;
  logic reset = 1, in_valid = 0, out_ready = 0;
  logic signed [7:0] in_pixel = 0;
  logic in_ready, out_valid, done;
  logic [1:0] out_row, out_col;
  logic signed [23:0] out_acc;
  logic signed [7:0] out_code;
  logic [7:0] out_activation, pooled;
  logic dense_ready, dense_valid, dense_done;
  logic [1:0] dense_row, dense_col;
  logic signed [23:0] dense_acc;
  logic signed [7:0] dense_code;
  logic [7:0] dense_activation, dense_pooled;

  cnn_conv3x3 dut (
    .clk(clk), .reset(reset), .in_valid(in_valid), .in_ready(in_ready),
    .in_pixel(in_pixel), .out_valid(out_valid), .out_ready(out_ready),
    .out_row(out_row), .out_col(out_col), .out_acc(out_acc), .out_code(out_code),
    .out_activation(out_activation), .done(done), .pooled(pooled)
  );
  // Nontrivial signed coefficients exercise products beyond +1,0,-1.
  cnn_conv3x3 #(
    .K00(8'sd1), .K01(8'sd2), .K02(8'sd3),
    .K10(-8'sd1), .K11(-8'sd2), .K12(-8'sd3),
    .K20(8'sd4), .K21(8'sd0), .K22(-8'sd4)
  ) dense (
    .clk(clk), .reset(reset), .in_valid(in_valid), .in_ready(dense_ready),
    .in_pixel(in_pixel), .out_valid(dense_valid), .out_ready(out_ready),
    .out_row(dense_row), .out_col(dense_col), .out_acc(dense_acc), .out_code(dense_code),
    .out_activation(dense_activation), .done(dense_done), .pooled(dense_pooled)
  );

  integer pixels [0:15];
  integer accepted = 0, received = 0, windows_seen = 0;
  integer frames_passed = 0, stalled_cycles = 0, expected_pool = 0, expected_dense_pool = 0;
  logic was_stalled = 0, window_was_stalled = 0;
  logic [43:0] held_output;
  logic [75:0] held_window;

  function automatic integer coefficient(input integer index, input integer dense_mode);
    begin
      if (!dense_mode) begin
        case (index % 3)
          0: coefficient = 1;
          1: coefficient = 0;
          default: coefficient = -1;
        endcase
      end else begin
        case (index)
          0: coefficient = 1; 1: coefficient = 2; 2: coefficient = 3;
          3: coefficient = -1; 4: coefficient = -2; 5: coefficient = -3;
          6: coefficient = 4; 7: coefficient = 0; default: coefficient = -4;
        endcase
      end
    end
  endfunction
  function automatic integer golden_acc(input integer r, input integer c, input integer dense_mode);
    integer kr, kc;
    begin
      golden_acc = 8;
      for (kr = 0; kr < 3; kr = kr + 1)
        for (kc = 0; kc < 3; kc = kc + 1)
          golden_acc = golden_acc + pixels[(r + kr)*4 + c + kc] * coefficient(kr*3 + kc, dense_mode);
    end
  endfunction
  function automatic integer golden_code(input integer value);
    integer floor_half;
    begin
      floor_half = value >>> 1;
      if ((value & 1) && (floor_half & 1)) floor_half = floor_half + 1;
      if (floor_half > 127) golden_code = 127;
      else if (floor_half < -128) golden_code = -128;
      else golden_code = floor_half;
    end
  endfunction
  function automatic integer golden_activation(input integer value);
    integer code;
    begin code = golden_code(value); golden_activation = code < 0 ? 0 : code; end
  endfunction
  task automatic check_window(input integer r, input integer c);
    begin
      if ($signed(dut.w00) !== pixels[r*4+c] || $signed(dut.w01) !== pixels[r*4+c+1] ||
          $signed(dut.w02) !== pixels[r*4+c+2] || $signed(dut.w10) !== pixels[(r+1)*4+c] ||
          $signed(dut.w11) !== pixels[(r+1)*4+c+1] || $signed(dut.w12) !== pixels[(r+1)*4+c+2] ||
          $signed(dut.w20) !== pixels[(r+2)*4+c] || $signed(dut.w21) !== pixels[(r+2)*4+c+1] ||
          $signed(dut.w22) !== pixels[(r+2)*4+c+2]) $fatal(1, "Window mismatch at %0d,%0d", r, c);
    end
  endtask

  always @(posedge clk) begin : scoreboard
    integer r, c, acc, dense_expected;
    if (reset) begin
      accepted = 0; received = 0; windows_seen = 0;
      expected_pool = 0; expected_dense_pool = 0;
      was_stalled = 0; window_was_stalled = 0;
    end else begin
      if (in_ready !== dense_ready || out_valid !== dense_valid || done !== dense_done)
        $fatal(1, "Kernel-independent handshake diverged");
      if (was_stalled && (!out_valid || {out_row,out_col,out_acc,out_code,out_activation} !== held_output))
        $fatal(1, "Output changed while stalled");
      if (window_was_stalled && (!dut.window_valid ||
          {dut.window_row,dut.window_col,dut.w00,dut.w01,dut.w02,dut.w10,dut.w11,dut.w12,dut.w20,dut.w21,dut.w22} !== held_window))
        $fatal(1, "Window changed while stalled");
      was_stalled = out_valid && !out_ready;
      window_was_stalled = dut.window_valid && !dut.window_ready;
      held_output = {out_row,out_col,out_acc,out_code,out_activation};
      held_window = {dut.window_row,dut.window_col,dut.w00,dut.w01,dut.w02,dut.w10,dut.w11,dut.w12,dut.w20,dut.w21,dut.w22};
      if (was_stalled) stalled_cycles = stalled_cycles + 1;
      if (in_valid && in_ready) begin
        if (accepted >= 16 || $signed(in_pixel) !== pixels[accepted]) $fatal(1, "Accepted input mismatch");
        accepted = accepted + 1;
      end
      if (dut.window_valid && dut.window_ready) begin
        if (windows_seen >= 4 || dut.window_row !== windows_seen/2 || dut.window_col !== windows_seen%2)
          $fatal(1, "Window order/count mismatch");
        check_window(windows_seen/2, windows_seen%2);
        windows_seen = windows_seen + 1;
      end
      if (out_valid && out_ready) begin
        if (received >= 4) $fatal(1, "Duplicate output");
        r = received/2; c = received%2;
        acc = golden_acc(r,c,0); dense_expected = golden_acc(r,c,1);
        if (out_row !== r || out_col !== c || $signed(out_acc) !== acc ||
            $signed(out_code) !== golden_code(acc) || out_activation !== golden_activation(acc))
          $fatal(1, "Default result mismatch at %0d,%0d got acc=%0d expected=%0d",r,c,out_acc,acc);
        if (dense_row !== r || dense_col !== c || $signed(dense_acc) !== dense_expected ||
            $signed(dense_code) !== golden_code(dense_expected) || dense_activation !== golden_activation(dense_expected))
          $fatal(1, "Dense signed-kernel result mismatch");
        if (golden_activation(acc) > expected_pool) expected_pool = golden_activation(acc);
        if (golden_activation(dense_expected) > expected_dense_pool) expected_dense_pool = golden_activation(dense_expected);
        received = received + 1;
      end
      if (done && (accepted != 16 || received != 4)) $fatal(1, "Premature completion");
    end
  end

  task automatic assert_reset;
    begin
      @(negedge clk); reset = 1; in_valid = 0; out_ready = 0;
      repeat (2) @(posedge clk);
      @(negedge clk); reset = 0;
      #1;
      if (out_valid || dense_valid || done || dense_done || !in_ready) $fatal(1, "Reset did not flush valid state");
    end
  endtask
  task automatic load_fixture(input integer fixture);
    integer i;
    begin
      for (i = 0; i < 16; i = i + 1) pixels[i] = i + 1;
      case (fixture)
        1: pixels[0] = 2; // Accumulator 3 -> nearest-even code 2.
        2: pixels[0] = -4; // Accumulator -3 -> nearest-even code -2 -> ReLU 0.
        3: for (i = 0; i < 16; i = i + 1) pixels[i] = (i%4 < 2) ? 127 : -128;
        4: for (i = 0; i < 16; i = i + 1) pixels[i] = (i%4 < 2) ? -128 : 127;
        default: begin end
      endcase
    end
  endtask
  task automatic run_frame(input integer fixture, input integer pauses);
    integer sent, cycles;
    logic hold_input;
    begin
      assert_reset(); load_fixture(fixture);
      sent = 0; cycles = 0; hold_input = 0;
      while (!done) begin
        @(negedge clk);
        out_ready = !pauses || (cycles%7 < 3);
        if (!hold_input) begin
          in_valid = sent < 16 && (!pauses || cycles%3 != 1);
          if (sent < 16) in_pixel = pixels[sent];
        end
        @(posedge clk);
        hold_input = in_valid && !in_ready;
        if (in_valid && in_ready) sent = sent + 1;
        #1; cycles = cycles + 1;
        if (cycles > 250) $fatal(1, "Frame timeout");
      end
      if (accepted != 16 || received != 4 || windows_seen != 4 || pooled !== expected_pool ||
          dense_pooled !== expected_dense_pool) $fatal(1, "Frame completion/count/pool mismatch");
      @(negedge clk); in_valid = 0; out_ready = 1;
      repeat (3) @(posedge clk);
      #1;
      if (out_valid || dense_valid || in_ready) $fatal(1, "Extra output or input after frame end");
      frames_passed = frames_passed + 1;
      $display("frame fixture=%0d pauses=%0d cycles=%0d outputs=4 pooled=%0d PASS",fixture,pauses,cycles,pooled);
    end
  endtask
  task automatic reset_with_pending_output;
    integer sent;
    begin
      assert_reset(); load_fixture(0); sent = 0;
      while (sent < 11) begin
        @(negedge clk); in_valid = 1; in_pixel = pixels[sent]; out_ready = 0;
        @(posedge clk); if (in_ready) sent = sent + 1;
      end
      @(negedge clk); in_valid = 0;
      repeat (3) @(posedge clk);
      #1;
      if (!out_valid || done) $fatal(1, "Pending-output reset fixture was not reached");
      assert_reset();
      $display("reset with pending output PASS");
    end
  endtask

  initial begin
    run_frame(0,0);
    run_frame(0,1);
    reset_with_pending_output();
    run_frame(1,1);
    run_frame(2,1);
    run_frame(3,1);
    run_frame(4,1);
    if (stalled_cycles == 0) $fatal(1, "No output stalls were exercised");
    $display("PASS frames=%0d windows=24 outputs=24 kernels=2 stalled_cycles=%0d reset_flush=1",frames_passed,stalled_cycles);
    $finish;
  end
  initial begin #20000; $fatal(1, "Global timeout"); end
endmodule