// Quantization-Aware Training Engine — Fake quantization + Straight-Through Estimator
module qat_engine #(parameter DW=16, QW=8, NC=64, SF=16)(
    input clk, rst_n, input train_mode, // 0=inference, 1=training
    input signed [DW-1:0] fwd_in, input fwd_valid, input [$clog2(NC)-1:0] ch,
    input [SF+8-1:0] scale[0:NC-1], input signed [QW-1:0] zp[0:NC-1],
    output reg signed [DW-1:0] fwd_out, output reg fwd_out_valid,
    // Backward path (STE gradient)
    input signed [DW-1:0] grad_in, input grad_valid,
    output reg signed [DW-1:0] grad_out, output reg grad_out_valid
);
    // Forward: fake quantize (quantize then dequantize)
    wire [SF+8-1:0] sc=scale[ch]; wire signed [QW-1:0] czp=zp[ch];
    wire signed [DW+SF+7:0] scaled=fwd_in*sc;
    wire signed [DW+SF+7:0] rounded=scaled+(1<<<(SF-1));
    wire signed [DW+7:0] iv=rounded[DW+SF+7:SF];
    wire signed [QW-1:0] qv=(iv>(1<<<(QW-1))-1)?((1<<<(QW-1))-1):(iv<(-(1<<<(QW-1))))?(-(1<<<(QW-1))):iv[QW-1:0];
    // Dequantize
    wire signed [DW+SF+7:0] dq_wide=(qv-czp)<<<SF;
    wire signed [DW-1:0] dq_val=dq_wide[DW+SF-1:SF];
    // STE gradient: pass gradient through if in range, else zero
    wire in_range=(iv>=( -(1<<<(QW-1))))&&(iv<=((1<<<(QW-1))-1));
    wire signed [DW-1:0] ste_grad=in_range?grad_in:'0;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin fwd_out<='0; fwd_out_valid<=0; grad_out<='0; grad_out_valid<=0; end
        else begin
            fwd_out<=train_mode?dq_val:fwd_in; fwd_out_valid<=fwd_valid;
            grad_out<=train_mode?ste_grad:grad_in; grad_out_valid<=grad_valid;
        end
    end
endmodule