// poly_mul_sync.v - Synchronous NTT-domain polynomial multiplier (BRAM storage)
//
// Computes pointwise (Karatsuba-like base-case) multiplication of two
// 256-coefficient NTT-domain polynomials.
//
// Operation flow:
//   IDLE -> LOAD (256x A+B pairs) -> RUN -> IDLE
//
// Coefficient storage moved from 256x12 LUT register arrays to four sd_bram
// banks (A_even/A_odd/B_even/B_odd, each 128x12): the base-case reads
// mem[2k] and mem[2k+1] together, so an even/odd split gives each bank a single
// read port. This removes the large multi-port LUT mux (~8k LUT) and the long
// BRAM->register-file routing that was the mlkem_clk critical path.
//
// RUN issues one base-case every two cycles. sd_bram reads are synchronous
// (1-cycle latency); the 2-cycle cadence absorbs it (address presented on the
// gap cycle, data consumed on the issue cycle). A one-cycle priming gap is
// inserted at LOAD->RUN so the first base-case's operands are ready.
//
// Interface (unchanged):
//   clk, rst_n     - Clock, active-low reset
//   coeff_a_in[11:0]- Polynomial A coefficient input
//   coeff_b_in[11:0]- Polynomial B coefficient input
//   valid_i         - Input valid
//   ready_o         - Ready to accept input (high in IDLE/LOAD)
//   coeff_out[11:0] - Result coefficient output
//   valid_o         - Output valid during RUN output phase
//   ready_i         - Output consumer ready

module poly_mul_sync (
    input         clk, rst_n,
    input  [11:0] coeff_a_in,
    input  [11:0] coeff_b_in,
    input         valid_i,
    output        ready_o,
    output [11:0] coeff_out,
    output        valid_o,
    input         ready_i
);

    localparam S_IDLE = 3'd0;
    localparam S_LOAD = 3'd1;
    localparam S_RUN  = 3'd2;

    reg [2:0] state, next_state;

    reg [7:0] load_cnt;       // 0..256 for loading 256 pairs
    reg [7:0] issue_cnt;      // 0..128, issued base-cases
    reg [7:0] out_cnt;        // 0..255, consumed output coefficients
    reg       issue_gap;      // spaces base-case requests to match c0/c1 output
    reg       out_phase;      // 0=c0, 1=c1
    reg       out_valid_r;

    reg [11:0] bc_a0_reg, bc_a1_reg;
    reg [11:0] bc_b0_reg, bc_b1_reg;
    reg [11:0] bc_zeta_reg;
    reg        bc_valid_reg;
    reg [11:0] c0_reg, c1_reg;

    // ---- Coefficient storage: even/odd banked sd_bram ----
    // A_even[k]=mem_A[2k], A_odd[k]=mem_A[2k+1]; same for B. Read index during
    // RUN is issue_k; write index during LOAD is load_cnt>>1 routed by load_cnt[0].
    wire [6:0] issue_k   = issue_cnt[6:0];       // 0..127 read index

    wire load_fire = ((state == S_IDLE) || (state == S_LOAD)) && valid_i && ready_o;
    wire [6:0] wr_local = load_cnt[7:1];          // load_cnt/2
    wire       wr_odd   = load_cnt[0];

    wire we_even = load_fire && !wr_odd;
    wire we_odd  = load_fire &&  wr_odd;

    wire [11:0] mem_a0, mem_a1, mem_b0, mem_b1;   // sd_bram registered read outputs
    sd_bram #(.W(12), .D(128), .A(7)) u_a_even (
        .clk(clk), .rd_addr(issue_k), .rd_data(mem_a0),
        .wr_en(we_even), .wr_addr(wr_local), .wr_data(coeff_a_in)
    );
    sd_bram #(.W(12), .D(128), .A(7)) u_a_odd (
        .clk(clk), .rd_addr(issue_k), .rd_data(mem_a1),
        .wr_en(we_odd), .wr_addr(wr_local), .wr_data(coeff_a_in)
    );
    sd_bram #(.W(12), .D(128), .A(7)) u_b_even (
        .clk(clk), .rd_addr(issue_k), .rd_data(mem_b0),
        .wr_en(we_even), .wr_addr(wr_local), .wr_data(coeff_b_in)
    );
    sd_bram #(.W(12), .D(128), .A(7)) u_b_odd (
        .clk(clk), .rd_addr(issue_k), .rd_data(mem_b1),
        .wr_en(we_odd), .wr_addr(wr_local), .wr_data(coeff_b_in)
    );

    wire [11:0] zeta;
    poly_mul_zeta_rom u_zeta (.addr(issue_k), .zeta(zeta));

    wire [11:0] bc_c0, bc_c1;
    wire        bc_vo;
    basecase_mul_pipe u_bc (
        .clk(clk), .rst_n(rst_n),
        .valid_i(bc_valid_reg),
        .a0(bc_a0_reg), .a1(bc_a1_reg),
        .b0(bc_b0_reg), .b1(bc_b1_reg),
        .zeta(bc_zeta_reg),
        .c0(bc_c0), .c1(bc_c1),
        .valid_o(bc_vo)
    );

    assign ready_o = (state == S_IDLE) || (state == S_LOAD);
    assign valid_o = out_valid_r;
    assign coeff_out = out_phase ? c1_reg : c0_reg;

    always @* begin
        next_state = state;
        case (state)
            S_IDLE:     if (valid_i && ready_o) next_state = S_LOAD;
            S_LOAD:     if (load_cnt >= 8'd255 && valid_i && ready_o)
                            next_state = S_RUN;
            S_RUN:      if (valid_o && ready_i && out_cnt == 8'd255)
                            next_state = S_IDLE;
            default:    next_state = S_IDLE;
        endcase
    end

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state    <= S_IDLE;
            load_cnt <= 8'd0;
            issue_cnt <= 8'd0;
            out_cnt <= 8'd0;
            issue_gap <= 1'b0;
            out_phase <= 1'b0;
            out_valid_r <= 1'b0;
            bc_a0_reg <= 12'd0; bc_a1_reg <= 12'd0;
            bc_b0_reg <= 12'd0; bc_b1_reg <= 12'd0;
            bc_zeta_reg <= 12'd0; bc_valid_reg <= 1'b0;
            c0_reg   <= 12'd0; c1_reg <= 12'd0;
        end else begin
            state <= next_state;
            bc_valid_reg <= 1'b0;

            // ---- LOAD phase (writes routed to even/odd banks combinationally) ----
            if (state == S_IDLE && valid_i && ready_o) begin
                load_cnt <= 8'd1;
                issue_cnt <= 8'd0;
                out_cnt <= 8'd0;
                issue_gap <= 1'b0;
                out_phase <= 1'b0;
                out_valid_r <= 1'b0;
            end

            if (state == S_LOAD && valid_i && ready_o) begin
                load_cnt <= load_cnt + 8'd1;
                if (load_cnt >= 8'd255) begin
                    issue_cnt <= 8'd0;
                    out_cnt <= 8'd0;
                    issue_gap <= 1'b1;   // priming gap: present issue_k=0 addr first
                    out_phase <= 1'b0;
                    out_valid_r <= 1'b0;
                end
            end

            // ---- COMPUTE phase ----
            // Issue one base-case every two cycles. On the gap cycle the sd_bram
            // address (issue_k) settles; on the issue cycle its registered output
            // (mem_a0/a1/b0/b1) is captured into the base-case input registers.
            if (state == S_RUN) begin
                if (issue_cnt < 8'd128) begin
                    if (!issue_gap) begin
                        bc_a0_reg <= mem_a0;
                        bc_a1_reg <= mem_a1;
                        bc_b0_reg <= mem_b0;
                        bc_b1_reg <= mem_b1;
                        bc_zeta_reg <= zeta;
                        bc_valid_reg <= 1'b1;
                        issue_cnt <= issue_cnt + 8'd1;
                        issue_gap <= 1'b1;
                    end else begin
                        issue_gap <= 1'b0;
                    end
                end

                if (out_valid_r && ready_i) begin
                    out_cnt <= out_cnt + 8'd1;
                    if (out_cnt == 8'd255) begin
                        out_valid_r <= 1'b0;
                        out_phase <= 1'b0;
                    end else if (!out_phase) begin
                        out_phase <= 1'b1;
                    end else begin
                        out_valid_r <= 1'b0;
                        out_phase <= 1'b0;
                    end
                end

                if (bc_vo) begin
                    c0_reg <= bc_c0;
                    c1_reg <= bc_c1;
                    out_valid_r <= 1'b1;
                    out_phase <= 1'b0;
                end
            end

            if (state == S_RUN && valid_o && ready_i && out_cnt == 8'd255) begin
                load_cnt <= 8'd0;
                issue_cnt <= 8'd0;
                out_cnt <= 8'd0;
                issue_gap <= 1'b0;
                out_phase <= 1'b0;
                out_valid_r <= 1'b0;
            end
        end
    end

endmodule