// Attention Mechanism Hardware — Scaled Dot-Product Attention
module attention_hw #(parameter DW=16, SEQ=64, DK=64, F=8)(
    input clk, rst_n, input en, input start,
    // Q, K, V matrices: SEQ × DK (row by row)
    input signed [DW-1:0] q_row [0:DK-1], input q_valid, input [6:0] q_idx,
    input signed [DW-1:0] k_row [0:DK-1], input k_valid, input [6:0] k_idx,
    input signed [DW-1:0] v_row [0:DK-1], input v_valid, input [6:0] v_idx,
    // Output
    output reg signed [DW-1:0] attn_out [0:DK-1],
    output reg [6:0] out_row, output reg out_valid, output reg done
);
    // Step 1: Score = Q × K^T / sqrt(dk)
    reg signed [DW*2-1:0] scores [0:SEQ-1]; // One row of scores
    reg signed [DW*2-1:0] score_sum;
    reg [6:0] q_r, k_r, d_r;
    reg [3:0] state;
    integer i;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin state<=0; q_r<=0; k_r<=0; d_r<=0; score_sum<='0; out_valid<=0; done<=0;
            for(i=0;i<SEQ;i++) scores[i]<='0; for(i=0;i<DK;i++) attn_out[i]<='0; out_row<=0; end
        else if(en) case(state)
          0: if(start) begin state<=1; q_r<=0; k_r<=0; end
          1: begin // Compute Q × K^T (dot product for one query row)
                score_sum <= score_sum + q_row[d_r] * k_row[d_r];
                d_r <= d_r + 1;
                if(d_r >= DK-1) begin d_r<=0;
                    scores[k_r] <= score_sum >>> 3; // /sqrt(64)=/8
                    score_sum <= '0; k_r <= k_r + 1;
                    if(k_r >= SEQ-1) begin k_r<=0; state<=2; end
                end
            end
          2: begin state<=3; // Softmax (simplified: find max, subtract, exp, normalize)
                // Simplified: just normalize scores to sum to 1<<F
                score_sum <= '0;
                for(i=0;i<SEQ;i++) score_sum <= score_sum + scores[i];
            end
          3: begin state<=4;
                // Apply attention weights to V
                for(i=0;i<SEQ;i++) begin
                    // out += softmax(score_i) * V_i
                    attn_out[d_r] <= attn_out[d_r] + ((scores[i] <<< F) / score_sum) * v_row[d_r]; // Simplified
                end
                d_r <= d_r + 1;
                if(d_r >= DK-1) begin d_r<=0; state<=4; end
            end
          4: begin out_valid<=1; out_row<=q_r; q_r<=q_r+1;
                if(q_r>=SEQ-1) state<=5; else state<=1;
                for(i=0;i<DK;i++) attn_out[i]<='0;
            end
          5: done<=1;
        endcase
    end
endmodule