// butterfly_unit_pipe.v - pipelined Cooley-Tukey / Gentleman-Sande butterfly
//
// Fixed-latency wrapper around barrett_mul_pipe. All arithmetic that used to
// sit in one NTT cycle is now split by registers.

module butterfly_unit_pipe (
    input         clk,
    input         rst_n,
    input         valid_i,
    input  [11:0] a,
    input  [11:0] b,
    input  [11:0] zeta,
    input         mode,       // 0 = forward NTT, 1 = inverse NTT
    output [11:0] a_out,
    output [11:0] b_out,
    output        valid_o
);
    localparam [11:0] Q12 = 12'd3329;
    localparam [12:0] Q13 = 13'd3329;

    // Plan A: register the inputs. Operands arrive combinationally from the
    // coeff_bram read; without this stage the path coeff_bram-read -> diff ->
    // first multiply stage sits in one cycle (the mlkem_clk critical path at
    // 100 MHz). Registering here splits it into (BRAM read -> reg) and
    // (reg -> diff -> multiply). Adds one pipeline cycle (absorbed by the NTT
    // writeback delay line).
    reg [11:0] a_r, b_r, zeta_r;
    reg        mode_ri, vld_r;

    wire [11:0] diff = (b_r >= a_r) ? (b_r - a_r) : (b_r - a_r + Q12);
    wire [11:0] mul_data = mode_ri ? diff : b_r;

    wire [11:0] mul_result;
    wire        mul_valid;
    barrett_mul_pipe u_barrett (
        .clk(clk),
        .rst_n(rst_n),
        .valid_i(vld_r),
        .a(zeta_r),
        .b(mul_data),
        .product(mul_result),
        .valid_o(mul_valid)
    );

    // a/b/mode delay lines (fed from the REGISTERED inputs a_r/b_r/mode_ri so
    // a_d6 aligns with mul_valid = valid_i delayed 8).
    reg [11:0] a_d0, a_d1, a_d2, a_d3, a_d4, a_d5, a_d6;
    reg [11:0] b_d0, b_d1, b_d2, b_d3, b_d4, b_d5, b_d6;
    reg        mode_d0, mode_d1, mode_d2, mode_d3, mode_d4, mode_d5, mode_d6;

    reg [11:0] a_out_r, b_out_r;
    reg        valid_r;

    wire [12:0] a_sum = {1'b0, a_d6} + {1'b0, (mode_d6 ? b_d6 : mul_result)};
    wire [11:0] a_mod = (a_sum >= Q13) ? (a_sum[11:0] - Q12) : a_sum[11:0];
    wire [11:0] b_sub = (a_d6 >= mul_result) ? (a_d6 - mul_result) :
                                                 (a_d6 - mul_result + Q12);

    assign a_out = a_out_r;
    assign b_out = b_out_r;
    assign valid_o = valid_r;

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            a_r <= 12'd0; b_r <= 12'd0; zeta_r <= 12'd0; mode_ri <= 1'b0; vld_r <= 1'b0;
            a_d0 <= 12'd0; a_d1 <= 12'd0; a_d2 <= 12'd0; a_d3 <= 12'd0; a_d4 <= 12'd0; a_d5 <= 12'd0; a_d6 <= 12'd0;
            b_d0 <= 12'd0; b_d1 <= 12'd0; b_d2 <= 12'd0; b_d3 <= 12'd0; b_d4 <= 12'd0; b_d5 <= 12'd0; b_d6 <= 12'd0;
            mode_d0 <= 1'b0; mode_d1 <= 1'b0; mode_d2 <= 1'b0; mode_d3 <= 1'b0; mode_d4 <= 1'b0; mode_d5 <= 1'b0; mode_d6 <= 1'b0;
            a_out_r <= 12'd0;
            b_out_r <= 12'd0;
            valid_r <= 1'b0;
        end else begin
            // input register stage
            a_r <= a; b_r <= b; zeta_r <= zeta; mode_ri <= mode; vld_r <= valid_i;

            a_d0 <= a_r;  a_d1 <= a_d0;    a_d2 <= a_d1;    a_d3 <= a_d2;    a_d4 <= a_d3;    a_d5 <= a_d4;    a_d6 <= a_d5;
            b_d0 <= b_r;  b_d1 <= b_d0;    b_d2 <= b_d1;    b_d3 <= b_d2;    b_d4 <= b_d3;    b_d5 <= b_d4;    b_d6 <= b_d5;
            mode_d0 <= mode_ri; mode_d1 <= mode_d0; mode_d2 <= mode_d1;
            mode_d3 <= mode_d2; mode_d4 <= mode_d3; mode_d5 <= mode_d4; mode_d6 <= mode_d5;

            valid_r <= mul_valid;
            if (mul_valid) begin
                a_out_r <= a_mod;
                b_out_r <= mode_d6 ? mul_result : b_sub;
            end
        end
    end

endmodule