module heron_fsm_tb;
  task assert;
    input condition;
    input [(20*8 - 1):0] message;
    begin
      if (condition !== 1'b1)
      begin
        $display("Assertion Failed: %s", message);
        $finish(2);
      end
    end
  endtask

  parameter CYCLE = 50;

  reg clk, reset;
  reg load;
  reg flag_z, flag_c, flag_s;

  wire[3:0] state;

  heron_fsm fsm(
    .clk(clk),
    .reset(reset),
    .load(load),
    .state(state),
    .flag_z(flag_z),
    .flag_c(flag_c),
    .flag_s(flag_s)
  );

  initial clk = 1'b0;
  always #(CYCLE/2) clk = ~clk;

  initial begin
    $dumpfile("fsm.vcd");
    $dumpvars;
  end

  initial begin
    reset = 1'b1;
    #(3*CYCLE)
    reset = 1'b0;
  end

  initial begin
    load = 0;
    #(5*CYCLE)
    load = 1;
    #(CYCLE)
    load = 0;

    assert(state == fsm.LD_N_1, "IDLE -> LD_N_1");
    #(CYCLE)
    assert(state == fsm.LD_N_2, "LD_N_1 -> LD_N_2");
    #(CYCLE)
    assert(state == fsm.I_GT_ZERO, "LD_N_2 -> I_GT_ZERO");

    $display("All tests passed!");
    $finish;
  end
endmodule