// butterfly_unit.v - Cooley-Tukey / Gentleman-Sande butterfly
//
// Computes one butterfly operation for NTT or inverse NTT.
//
// Parameters:
//   Q = 3329 (prime modulus)
//
// Forward NTT (mode=0):
//   t = zeta * b mod Q (via barrett_mul)
//   a_out = (a + t) mod Q
//   b_out = (a - t) mod Q
//
// Inverse NTT (mode=1):
//   a_out = (a + b) mod Q
//   diff = (b - a) mod Q   (handled as: if b >= a: b-a; else: b-a+Q)
//   b_out = zeta * diff mod Q (via barrett_mul)
//
// Pure combinational.

module butterfly_unit (
    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
);
    localparam Q = 3329;

    // Barrett modular multiplication: zeta * mul_data mod Q
    wire [11:0] mul_data;     // what to multiply with zeta
    wire [11:0] mul_result;   // result of barrett_mul(zeta, mul_data)

    // Forward:  mul_data = b,          t = zeta * b mod Q
    // Inverse:  mul_data = (b - a) mod Q positive
    assign mul_data = (mode == 1'b0) ? b : ((b >= a) ? (b - a) : (b - a + Q));

    barrett_mul u_barrett (
        .a       (zeta),
        .b       (mul_data),
        .product (mul_result)
    );

    // ---- a_out computation ----
    // Forward:  a_out = (a + t) mod Q
    // Inverse:  a_out = (a + b) mod Q
    wire [12:0] a_sum;
    wire [11:0] add_val;
    assign add_val = (mode == 1'b0) ? mul_result : b;
    assign a_sum = {1'b0, a} + {1'b0, add_val};
    // a_sum - Q produces 13-bit result; we only need lower 12 bits since
    // a_sum >= Q guarantees the result < Q < 2^12
    wire [11:0] a_sub_12;
    assign a_sub_12 = a_sum[11:0] - Q[11:0];
    wire [11:0] a_result = (a_sum >= Q) ? a_sub_12 : a_sum[11:0];
    assign a_out = a_result;

    // ---- b_out computation ----
    // Forward:  b_out = (a - t) mod Q  → if a >= t: a-t; else: a-t+Q
    // Inverse:  b_out = t (mul_result, which is zeta * (b-a) mod Q)
    wire [11:0] sub_val;
    assign sub_val = mul_result;
    wire [11:0] sub_result;
    assign sub_result = (a >= sub_val) ? (a - sub_val) : (a - sub_val + Q);
    assign b_out = (mode == 1'b0) ? sub_result : mul_result;

endmodule