`timescale 1ns / 1ps
`include "cam_params.svh"

module noise_mask_bernoulli #(
    parameter int HASH_BITS      = `HASH_BITS,
    parameter int PROB_BITS      = 8,
    parameter int PRNG_WORDS     = 2,
    parameter int BITS_PER_CYCLE = 32,
    parameter logic [63:0] SEED  = 64'hD1B5_4A32_9E37_79B9
) (
    input  logic                       clk,
    input  logic                       rst_n,
    input  logic                       start_i,
    input  logic [PROB_BITS-1:0]       threshold_i,
    output logic                       busy_o,
    output logic                       done_o,
    output logic [HASH_BITS-1:0]       mask_o
);

    localparam int RANDOM_BITS = PRNG_WORDS * 128;
    localparam int BLOCKS = HASH_BITS / BITS_PER_CYCLE;
    localparam int BLOCK_INDEX_BITS = (BLOCKS <= 1) ? 1 : $clog2(BLOCKS);
    localparam logic [BLOCK_INDEX_BITS-1:0] LAST_BLOCK_IDX = BLOCK_INDEX_BITS'(BLOCKS - 1);

    typedef enum logic [1:0] {
        STATE_IDLE,
        STATE_PRIME,
        STATE_RUN
    } state_t;

    state_t state_q;
    logic [BLOCK_INDEX_BITS-1:0] block_idx_q;
    logic [PROB_BITS-1:0] threshold_q;
    logic [HASH_BITS-1:0] mask_q;
    logic prng_en;
    logic [RANDOM_BITS-1:0] random_bits;
    logic [BITS_PER_CYCLE-1:0] mask_slice;

    assign busy_o = (state_q != STATE_IDLE);
    assign mask_o = mask_q;
    assign prng_en = (state_q == STATE_PRIME) ||
                     ((state_q == STATE_RUN) && (block_idx_q != LAST_BLOCK_IDX));

    `ifndef SYNTHESIS
    initial begin
        if (HASH_BITS <= 0) $fatal(1, "HASH_BITS must be > 0");
        if (PROB_BITS <= 0) $fatal(1, "PROB_BITS must be > 0");
        if (PRNG_WORDS <= 0) $fatal(1, "PRNG_WORDS must be > 0");
        if (BITS_PER_CYCLE <= 0) $fatal(1, "BITS_PER_CYCLE must be > 0");
        if (HASH_BITS % BITS_PER_CYCLE != 0) $fatal(1, "HASH_BITS must be divisible by BITS_PER_CYCLE");
        if (RANDOM_BITS != BITS_PER_CYCLE * PROB_BITS) $fatal(1, "PRNG_WORDS*128 must equal BITS_PER_CYCLE*PROB_BITS");
        if (SEED == 64'h0) $fatal(1, "SEED must be non-zero");
        for (int w = 0; w < PRNG_WORDS; w++) begin
            if ({SEED, SEED ^ (64'(w+1))} == 128'h0) $fatal(1, "PRNG seed must be non-zero");
            for (int w2 = w+1; w2 < PRNG_WORDS; w2++) begin
                if ({SEED, SEED ^ (64'(w+1))} == {SEED, SEED ^ (64'(w2+1))})
                    $fatal(1, "PRNG seeds must differ");
            end
        end
    end
    `endif

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_q <= STATE_IDLE;
            block_idx_q <= '0;
            threshold_q <= '0;
            mask_q <= '0;
            done_o <= 1'b0;
        end else begin
            done_o <= 1'b0;

            unique case (state_q)
                STATE_IDLE: begin
                    block_idx_q <= '0;
                    if (start_i) begin
                        threshold_q <= threshold_i;
                        mask_q <= '0;
                        state_q <= STATE_PRIME;
                    end
                end

                STATE_PRIME: begin
                    state_q <= STATE_RUN;
                end

                STATE_RUN: begin
                    mask_q[block_idx_q * BITS_PER_CYCLE +: BITS_PER_CYCLE] <= mask_slice;

                    if (block_idx_q == LAST_BLOCK_IDX) begin
                        state_q <= STATE_IDLE;
                        block_idx_q <= '0;
                        done_o <= 1'b1;
                    end else begin
                        block_idx_q <= block_idx_q + 1'b1;
                    end
                end

                default: begin
                    state_q <= STATE_IDLE;
                    block_idx_q <= '0;
                end
            endcase
        end
    end

    generate
        for (genvar w = 0; w < PRNG_WORDS; w++) begin : gen_prng
            localparam logic [127:0] PRNG_SEED = {SEED, SEED ^ (64'(w+1))};
            random128 prng_inst (
                .clk    (clk),
                .rst_n  (rst_n),
                .enable (prng_en),
                .seed   (PRNG_SEED),
                .out    (random_bits[w*128 +: 128])
            );
        end
    endgenerate

    generate
        for (genvar i = 0; i < BITS_PER_CYCLE; i++) begin : gen_bernoulli
            assign mask_slice[i] = (random_bits[i*PROB_BITS +: PROB_BITS] < threshold_q);
        end
    endgenerate

endmodule