/* verilator lint_off BLKLOOPINIT */
// Knowledge Distillation Hardware — Teacher-Student loss combiner
module distill_hw #(parameter DW=16, NCLASS=1000, TEMP_W=8, ALPHA_W=8)(
    input clk, rst_n, input en,
    input signed [DW-1:0] student_logits [0:NCLASS-1], input s_valid,
    input signed [DW-1:0] teacher_logits [0:NCLASS-1], input t_valid,
    input [TEMP_W-1:0] temperature, // Softmax temperature
    input [ALPHA_W-1:0] alpha,       // Distillation loss weight (0.0-1.0 in fixed point)
    input [3:0] true_label,          // Ground truth label for hard loss
    // Output: combined loss gradient
    output reg signed [DW-1:0] grad_out [0:NCLASS-1], output reg grad_valid,
    output reg signed [DW-1:0] soft_loss, hard_loss, total_loss
);
    // Softmax approximation: exp(x/T) / sum(exp(x/T))
    // Using piecewise linear approximation of exp
    reg signed [DW-1:0] s_soft [0:NCLASS-1];
    reg signed [DW-1:0] t_soft [0:NCLASS-1];
    reg [15:0] s_sum, t_sum;
    integer ci;
    always_ff @(posedge clk or negedge rst_n) begin
        if(!rst_n) begin for(ci=0;ci<NCLASS;ci++) begin s_soft[ci]<='0; t_soft[ci]<='0; grad_out[ci]<='0; end
            s_sum<=0; t_sum<=0; soft_loss<='0; hard_loss<='0; total_loss<='0; grad_valid<=0; end
        else if(en && s_valid && t_valid) begin
            s_sum<=0; t_sum<=0;
            for(ci=0;ci<NCLASS;ci++) begin
                // Simplified: scaled logits as proxy for softmax
                s_soft[ci] <= student_logits[ci] >>> temperature;
                t_soft[ci] <= teacher_logits[ci] >>> temperature;
                s_sum <= s_sum + student_logits[ci];
                t_sum <= t_sum + teacher_logits[ci];
            end
            // KL divergence gradient: alpha * (s_soft - t_soft)
            // Hard loss gradient: (1-alpha) * (s_soft - one_hot)
            for(ci=0;ci<NCLASS;ci++) begin
                soft_loss <= t_soft[ci] - s_soft[ci]; // Simplified KL gradient
                if(ci == true_label) hard_loss <= s_soft[ci] - 1; // CE gradient for true class
                else hard_loss <= s_soft[ci];
                grad_out[ci] <= (alpha * (t_soft[ci] - s_soft[ci]) + (1 - alpha) * hard_loss) >>> ALPHA_W;
            end
            total_loss <= soft_loss + hard_loss;
            grad_valid <= 1;
        end else grad_valid <= 0;
    end
endmodule