// comp_decomp_sync.v - ML-KEM coefficient compression/decompression
//
// Streaming valid/ready interface. The arithmetic is fixed-latency pipelined
// so compression does not infer a long combinational divider/reducer.
// mode=0: compress   — round((2^d * x) / Q) mod 2^d
// mode=1: decompress — round((Q * x) / 2^d)
// 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 (always < Q)

`include "sync_rtl/common/defines.vh"

(* use_dsp = "yes" *)
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
);

    localparam [24:0] Q25 = 25'd`Q;
    localparam [12:0] K13 = 13'd5039;   // floor(2^24 / 3329)

    reg [11:0] data_r;
    reg        valid_r;

    reg        mode_s0;
    reg [4:0]  d_s0;
    reg [11:0] coeff_s0;
    reg        valid_s0;

    reg        mode_s1;
    reg [4:0]  d_s1;
    reg [23:0] comp_dividend_s1;
    reg [23:0] decomp_rounded_s1;
    reg        valid_s1;

    reg        mode_s2;
    reg [4:0]  d_s2;
    reg [23:0] comp_dividend_s2;
    reg [36:0] comp_qprod_est_s2;
    reg [23:0] decomp_rounded_s2;
    reg        valid_s2;

    reg        mode_s3;
    reg [4:0]  d_s3;
    reg [23:0] comp_dividend_s3;
    reg [12:0] comp_q_est_s3;
    reg [11:0] decomp_result_s3;
    reg        valid_s3;

    reg        mode_s4;
    reg [4:0]  d_s4;
    reg [23:0] comp_dividend_s4;
    reg [12:0] comp_q_est_s4;
    reg [24:0] comp_q_mul_q_s4;
    reg [11:0] decomp_result_s4;
    reg        valid_s4;

    reg        mode_s5;
    reg [4:0]  d_s5;
    reg [12:0] comp_q_est_s5;
    reg [24:0] comp_rem_s5;
    reg [11:0] decomp_result_s5;
    reg        valid_s5;

    wire fire_i = valid_i && ready_o;

    wire pipe_busy = valid_s0 | valid_s1 | valid_s2 | valid_s3 | valid_s4 | valid_s5;

    // Accept a new command only when the pipeline and output slot are both free.
    // Current users issue one coefficient at a time and wait for valid_o.
    assign ready_o   = !pipe_busy && (!valid_r || ready_i);
    assign coeff_out = data_r;
    assign valid_o   = valid_r;

    wire [23:0] coeff_ext_s0 = {12'd0, coeff_s0};
    wire [23:0] comp_dividend_next = (coeff_ext_s0 << d_s0) + 24'd1664;
    wire [23:0] decomp_product_next = coeff_s0 * 12'd`Q;
    wire [23:0] decomp_round_next =
        decomp_product_next + ((d_s0 == 5'd0) ? 24'd0 : (24'd1 << (d_s0 - 5'd1)));

    wire [12:0] comp_q_corr1_s5 =
        comp_q_est_s5 + ((comp_rem_s5 >= Q25) ? 13'd1 : 13'd0);
    wire [24:0] comp_rem_corr1_s5 =
        (comp_rem_s5 >= Q25) ? (comp_rem_s5 - Q25) : comp_rem_s5;
    wire [12:0] comp_q_corr2_s5 =
        comp_q_corr1_s5 + ((comp_rem_corr1_s5 >= Q25) ? 13'd1 : 13'd0);
    wire [11:0] comp_mask_s5 =
        (d_s5 == 5'd0) ? 12'd0 : (12'd1 << d_s5) - 12'd1;
    wire [11:0] comp_result_s5 = comp_q_corr2_s5[11:0] & comp_mask_s5;

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            data_r <= 12'd0;
            valid_r <= 1'b0;

            mode_s0 <= 1'b0;
            d_s0 <= 5'd0;
            coeff_s0 <= 12'd0;
            valid_s0 <= 1'b0;

            mode_s1 <= 1'b0;
            d_s1 <= 5'd0;
            comp_dividend_s1 <= 24'd0;
            decomp_rounded_s1 <= 24'd0;
            valid_s1 <= 1'b0;

            mode_s2 <= 1'b0;
            d_s2 <= 5'd0;
            comp_dividend_s2 <= 24'd0;
            comp_qprod_est_s2 <= 37'd0;
            decomp_rounded_s2 <= 24'd0;
            valid_s2 <= 1'b0;

            mode_s3 <= 1'b0;
            d_s3 <= 5'd0;
            comp_dividend_s3 <= 24'd0;
            comp_q_est_s3 <= 13'd0;
            decomp_result_s3 <= 12'd0;
            valid_s3 <= 1'b0;

            mode_s4 <= 1'b0;
            d_s4 <= 5'd0;
            comp_dividend_s4 <= 24'd0;
            comp_q_est_s4 <= 13'd0;
            comp_q_mul_q_s4 <= 25'd0;
            decomp_result_s4 <= 12'd0;
            valid_s4 <= 1'b0;

            mode_s5 <= 1'b0;
            d_s5 <= 5'd0;
            comp_q_est_s5 <= 13'd0;
            comp_rem_s5 <= 25'd0;
            decomp_result_s5 <= 12'd0;
            valid_s5 <= 1'b0;
        end else begin
            if (valid_r && ready_i)
                valid_r <= 1'b0;

            mode_s0 <= mode;
            d_s0 <= d;
            coeff_s0 <= coeff_in;
            valid_s0 <= fire_i;

            mode_s1 <= mode_s0;
            d_s1 <= d_s0;
            comp_dividend_s1 <= comp_dividend_next;
            decomp_rounded_s1 <= decomp_round_next;
            valid_s1 <= valid_s0;

            mode_s2 <= mode_s1;
            d_s2 <= d_s1;
            comp_dividend_s2 <= comp_dividend_s1;
            comp_qprod_est_s2 <= {13'd0, comp_dividend_s1} * K13;
            decomp_rounded_s2 <= decomp_rounded_s1;
            valid_s2 <= valid_s1;

            mode_s3 <= mode_s2;
            d_s3 <= d_s2;
            comp_dividend_s3 <= comp_dividend_s2;
            comp_q_est_s3 <= comp_qprod_est_s2[36:24];
            decomp_result_s3 <= decomp_rounded_s2 >> d_s2;
            valid_s3 <= valid_s2;

            mode_s4 <= mode_s3;
            d_s4 <= d_s3;
            comp_dividend_s4 <= comp_dividend_s3;
            comp_q_est_s4 <= comp_q_est_s3;
            comp_q_mul_q_s4 <= comp_q_est_s3 * Q25;
            decomp_result_s4 <= decomp_result_s3;
            valid_s4 <= valid_s3;

            mode_s5 <= mode_s4;
            d_s5 <= d_s4;
            comp_q_est_s5 <= comp_q_est_s4;
            comp_rem_s5 <= {1'b0, comp_dividend_s4} - comp_q_mul_q_s4;
            decomp_result_s5 <= decomp_result_s4;
            valid_s5 <= valid_s4;

            if (valid_s5) begin
                data_r <= mode_s5 ? decomp_result_s5 : comp_result_s5;
                valid_r <= 1'b1;
            end
        end
    end

endmodule