module last_heap_stage #(
    parameter VALUE_WIDTH	= 32,
    parameter KEY_WIDTH 	= 32
) (
    input wire 			    ap_clk,
    input wire 			    ap_reset,

    input wire [KEY_WIDTH - 1:0]    key_test_in,    // input key to be tested
    input wire [KEY_WIDTH - 1:0]    key_in,         // input key from prev. heap stage
    input wire [VALUE_WIDTH - 1:0]  value_in,       // input value from prev. heap stage
    input wire 			    kv_in_valid,    // input key,value pair to be tested is ready
    input wire 			    heap_read,      // heap read request, from AXI-lite mechanism in CMS.

//    output wire [KEY_WIDTH - 1:0]   key_test_out,   // key to be tested, out to next heap stage
    output wire [KEY_WIDTH - 1:0]   key_out,        // current stage key, out to next heap stage
    output wire [VALUE_WIDTH - 1:0] value_out,      // current stage value, out to next heap stage
    output wire 		    kv_out_valid    // current key,value pair is valid, out to next heap stage
);

//------------------------heap stage fsm-------------------
localparam
  K_TEST      = 2'd0,
  V_TEST      = 2'd1,
  KV_WRITE    = 2'd2,
  PASSTHRU    = 2'd3;

// heap stage data structures:	  
// Option 1: Using regs:
//  reg [KEY_WIDTH - 1:0] 	    key_test;
  reg 				    kv_valid;
  reg 				    stage_valid;
  reg [KEY_WIDTH - 1:0] 	    key;
  reg [VALUE_WIDTH - 1:0] 	    value;
//  reg [KEY_WIDTH - 1:0] 	    key_out_r;
//  reg [VALUE_WIDTH - 1:0] 	    value_out_r;
  

  reg [1:0] 			    hstate;
  wire 				    key_is_tested;
  wire 				    heap_entry_is_bubbled;

  assign kv_out_valid = kv_valid;
//  assign key_test_out = key_test;
  assign key_out = key;
  assign value_out = value;
			
  always @(posedge ap_clk) begin
    
    if (ap_reset) begin
      key <= 0;
      value <= 0;
      kv_valid <= 0;
      hstate <= K_TEST;
      stage_valid <= 0;
    end
    
    else begin

      // Heap stage State Machine
      if (hstate == K_TEST) begin
	kv_valid <= 1'b0;
	
	if (heap_read) begin
	  hstate = PASSTHRU;
	  
	  // Reading thru AXI-lite: 
	  // heap stages are bubbled, two stages per each read request.
	  // The second bubble in PASSTHRU state   
	  //	    key_out_r <= key;
	  //	    value_out_r <= value;
	  key <= key_in;	
	  value <= value_in;
	end	    

	else if (kv_in_valid) begin
	  if (stage_valid & (key_test_in == key) | !stage_valid) begin
	    hstate = KV_WRITE;
	  end
	  else begin
	    hstate = V_TEST;
	  end
	end
      end
	
      else if (hstate == KV_WRITE) begin
	// The incoming KV pair is loaded into current stage if:
	// 1. Curent stage was empty (!stage_valid), or
	// 2. The incoming K matches the existing
	// In both cases, nothing is passed down to next stage
	hstate = K_TEST;
	stage_valid <= 1'b1;
	
	// Replace current kv with the incoming pair. 
	key <= key_in;	
	value <= value_in;
	
	// Prev. kv values are NOT forwarded to next heap stage, but still the output kv registers are also written, just to keep its write_enable signal common to all the registers in a heap stage
	//	  key_out_r <= key;
	//	  value_out_r <= value;
      end
      
      else if (hstate == V_TEST) begin
	hstate = K_TEST;
	kv_valid <= 1'b1;
	//	  key_test <= key_test_in;	
	
	if (value_in > value) begin
	  // The new value is greater than current value.
	  // Replace current kv with the incoming pair and forward prev. kv values to next heap stage 
	  key <= key_in;	
	  value <= value_in;
	  //	  key_out_r <= key;
	  //	  value_out_r <= value;
	end
	else begin
	  // The new value is not greater than current value.
	  // Current kv pair is maintained. Incoming pair is forwarded to next heap stage 
	  //	  key_out_r <= key_in;
	  //	  value_out_r <= value_in;
	end
      end
	  
      else if (hstate == PASSTHRU) begin
	hstate = K_TEST;
	stage_valid <= 0;
	
	// Reading thru AXI-lite: 
	// heap stages are bubbled, two stages per each read request.
	// The first bubble was done in prev. K_TEST state   
	//	  key_out_r <= key;
	//	  value_out_r <= value;
	key <= key_in;	
	value <= value_in;
      end	
	
      else
	hstate = K_TEST;
      
    end // else (if (ap_reset))
  end //always @(posedge ap_clk)

endmodule