// poly_mul_sync.v - Synchronous NTT-domain polynomial multiplier
//
// Computes pointwise (Karatsuba-like base-case) multiplication of two
// 256-coefficient NTT-domain polynomials.
//
// Operation flow:
//   IDLE -> LOAD (256x A+B pairs) -> RUN -> IDLE
//
// RUN issues one base-case multiply every two cycles, matching the two-cycle
// c0/c1 output bandwidth after the basecase pipeline fills.
//
// The LOAD phase accepts both A and B coefficients simultaneously
// (one pair per cycle) on coeff_a_in/coeff_b_in.
//
// The COMPUTE phase outputs the 256 result coefficients one per cycle
// via valid/ready handshake on coeff_out/valid_o/ready_i.
//
// Memory: 256×12-bit register arrays for A and B coefficients.
//
// Interface:
//   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
);

    // State definitions
    localparam S_IDLE = 3'd0;
    localparam S_LOAD = 3'd1;
    localparam S_RUN  = 3'd2;

    reg [2:0] state, next_state;

    // Coefficient storage (register arrays)
    reg [11:0] mem_A [0:255];
    reg [11:0] mem_B [0:255];

    // Counters
    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;

    // Registered basecase_mul inputs/results
    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;

    // Combinational read signals for the next base-case request.
    wire [6:0] issue_k   = issue_cnt[6:0];
    wire [7:0] addr_even = {issue_k, 1'b0};   // issue_k * 2
    wire [7:0] addr_odd  = {issue_k, 1'b1};   // issue_k * 2 + 1
    wire [11:0] mem_a0 = mem_A[addr_even];
    wire [11:0] mem_a1 = mem_A[addr_odd];
    wire [11:0] mem_b0 = mem_B[addr_even];
    wire [11:0] mem_b1 = mem_B[addr_odd];

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

    // Pipelined basecase multiply. Requests are issued every two cycles; inputs
    // are registered locally so issue_cnt does not directly drive DSP input muxes.
    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)
    );

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

    // State transition logic (combinational)
    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

    // Sequential logic
    integer i;
    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;
            for (i = 0; i < 256; i = i + 1) begin
                mem_A[i] <= 12'd0;
                mem_B[i] <= 12'd0;
            end
        end else begin
            state <= next_state;
            bc_valid_reg <= 1'b0;

            // ---- LOAD phase ----
            // First coefficient captured on IDLE → LOAD transition
            if (state == S_IDLE && valid_i && ready_o) begin
                mem_A[0] <= coeff_a_in;
                mem_B[0] <= coeff_b_in;
                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

            // Subsequent coefficients in LOAD state
            if (state == S_LOAD && valid_i && ready_o) begin
                mem_A[load_cnt] <= coeff_a_in;
                mem_B[load_cnt] <= coeff_b_in;
                load_cnt <= load_cnt + 8'd1;
                if (load_cnt >= 8'd255) begin
                    issue_cnt <= 8'd0;
                    out_cnt <= 8'd0;
                    issue_gap <= 1'b0;
                    out_phase <= 1'b0;
                    out_valid_r <= 1'b0;
                end
            end

            // ---- COMPUTE phase ----
            // RUN: cut issue_cnt -> memory mux -> basecase DSP with local input
            // regs, then launch requests every other cycle. The two-cycle issue
            // interval matches the c0/c1 output bandwidth, so no deep result FIFO
            // is needed for this always-ready system.
            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

            // ---- final output consumed: reset counters before accepting a new op ----
            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