// comp_decomp_sync.v - ML-KEM coefficient compression/decompression
//
// Streaming: one coefficient per cycle through pipeline_reg.
// mode=0: compress   - round((2^d * x) / Q) mod 2^d
// mode=1: decompress - round((Q * x) / 2^d) mod Q
// Uses round-half-up (round(2.5)=3, round(3.5)=4).
// Integer arithmetic:
//   compress:   (x * 2^d + Q/2) / Q,  lower d bits as result
//   decompress: (x * Q + 2^(d-1)) / 2^d, result mod Q
//
// Division by Q=3329 replaced with Barrett multiplication + correction:
//   floor(x / Q) ≈ (x * 5039) >> 24   (5039 = floor(2^24 / 3329))
//   Correction: if x - q_approx * Q >= Q, add 1 to quotient.
//   Q*5039 = 16,774,831 < 2^24, so Barrett may under-estimate by 1.
//   Q = 3329 = 2048+512+256+1, so q*Q uses 3 shifts + add (no 2nd multiplier).

`include "sync_rtl/common/defines.vh"

(* use_dsp = "no" *)
module comp_decomp_sync (
    input         clk,
    input         rst_n,
    input  [11:0] coeff_in,
    input  [4:0]  d,
    input         mode,          // 0=compress, 1=decompress
    input         valid_i,
    output        ready_o,
    output [11:0] coeff_out,
    output        valid_o,
    input         ready_i
);

    wire [11:0] two_pow_d = 12'd1 << d;

    // product: 12b*12b = max 2048*3328 = 6,815,744 → 23b, pad to 24b
    wire [23:0] product   = mode ? {12'b0, coeff_in} * {12'b0, 12'(`Q)}
                                 : {12'b0, coeff_in} * {12'b0, two_pow_d};

    wire [11:0] round_off = mode ? (two_pow_d >> 1) : 12'd1664;

    wire [23:0] dividend  = product + {12'b0, round_off};

    // ===========================================================
    // Barrett: floor(x / Q) ≈ (x * 5039) >> 24
    //   M = floor(2^24 / Q) = 5039
    //   M*Q = 16,774,831 < 2^24 → floor(x/Q) may be 1 too low
    //   Correction: remainder = x - q_approx*Q; if rem >= Q, q→q+1
    //   Q*5039 fails at: k*2385 >= 2^24 → k >= 7033; max k ≈ 2047 → safe
    // ===========================================================
    wire [36:0] barrett_prod = dividend * 13'd5039;            // 24×13=37b
    wire [11:0] q_approx     = barrett_prod[35:24];             // quotient estimate
    // q*Q = q*2048 + q*512 + q*256 + q  (3329 = 2048+512+256+1)
    wire [23:0] q_times_Q    = ({q_approx, 11'd0})
                             + ({q_approx, 9'd0})
                             + ({q_approx, 8'd0})
                             + q_approx;
    wire [24:0] delta        = {1'b0, dividend} - {1'b0, q_times_Q}; // 25b
    wire [11:0] comp_q       = q_approx + (delta >= {13'b0, 12'(`Q)});
    // lower d bits = compressed result
    wire [11:0] comp_res     = comp_q & ({12'b0, two_pow_d} - 1'b1);

    // ===========================================================
    // Decompress: dividend >> d (free), then % Q via Barrett
    // ===========================================================
    wire [23:0] dec_sh       = dividend >> d;
    wire [36:0] dm_prod      = dec_sh * 13'd5039;
    wire [11:0] dm_qe        = dm_prod[35:24];
    wire [23:0] dm_qQ        = ({dm_qe, 11'd0})
                             + ({dm_qe, 9'd0})
                             + ({dm_qe, 8'd0})
                             + dm_qe;
    wire [24:0] dm_del       = {1'b0, dec_sh} - {1'b0, dm_qQ};
    // divide by Q, take remainder: r = x - q*Q
    wire [12:0] dm_q         = dm_qe + (dm_del >= {13'b0, 12'(`Q)}); // corr quotient
    // actual remainder = dec_sh - dm_q*Q
    wire [23:0] dm_corrQ     = ({dm_q, 11'd0}) + ({dm_q, 9'd0}) + ({dm_q, 8'd0}) + dm_q;
    wire [24:0] dm_rem       = {1'b0, dec_sh} - {1'b0, dm_corrQ};
    // dm_rem < Q*2 (Barrett gives exact quotient after correction), one cond-sub
    wire [11:0] dec_mod      = (dm_rem >= {13'b0, 12'(`Q)})
                             ? (dm_rem[11:0] - 12'(`Q))
                             : dm_rem[11:0];

    wire [11:0] mod_result   = mode ? dec_mod : comp_res;

    pipeline_reg #(.DW(12)) u_pipe (
        .clk    (clk),
        .rst_n  (rst_n),
        .data_i (mod_result),
        .valid_i(valid_i),
        .ready_o(ready_o),
        .data_o (coeff_out),
        .valid_o(valid_o),
        .ready_i(ready_i)
    );

endmodule