// ntt_core.v - NTT core with individual coefficient registers (2 butterflies)
//
// Uses 256 individual 12-bit registers and TWO deeply pipelined butterfly paths
// running in lock-step (R1: dual-butterfly parallelism). Each issue cycle emits
// two butterflies from the SAME group at positions (p, p+1); since ML-KEM's NTT
// has 7 layers with layer_len in {128..2} (always even, never 1), the pair always
// shares one zeta and touches four distinct coefficient indices, so a single
// zeta ROM suffices and there is no intra-group read/write conflict.
//   address/operand/zeta register -> 2x pipelined Barrett butterfly -> writeback
// Each NTT layer issues 2 butterflies/cycle (64 cycles/layer vs 128), then the
// core drains pending writebacks before starting the next dependent layer.
// In inverse mode, final x3303 output scaling also uses a pipelined Barrett
// multiplier so the output path does not reintroduce a combinational reducer.

module ntt_core (
    input         clk, rst_n,
    input  [11:0] coeff_in,
    input         valid_i,
    output        ready_o,
    input         mode,
    output [11:0] coeff_out,
    output        valid_o,
    input         ready_i,
    output        done_o
);
    localparam N = 256, LAYERS = 7, DW = 12;

    reg [DW-1:0] cr [0:N-1];
    integer ci;

    localparam S_IDLE      = 4'd0;
    localparam S_LOAD      = 4'd1;
    localparam S_CMP_RUN   = 4'd2;
    localparam S_CMP_DRAIN = 4'd3;
    localparam S_OUT_PREP  = 4'd7;
    localparam S_OUTPUT    = 4'd8;
    localparam S_OUT_SCALE = 4'd9;
    localparam S_DONE      = 4'd10;

    reg [3:0] state, next_state;

    reg [7:0] load_cnt;
    reg [7:0] out_cnt;
    reg [8:0] scale_issue_cnt;
    reg [8:0] scale_emit_cnt;
    reg [7:0] j, start, layer_len;
    reg [6:0] zeta_idx;
    reg [2:0] layer;
    reg       bf_done;
    reg       mode_r;
    reg       layer_issue_done;
    reg [5:0] bf_inflight;    // pending issue cycles (pairs) in flight

    // Two writeback address delay lines (lane0: j, lane1: j+1), each +len.
    reg [7:0] a0_d0,a0_d1,a0_d2,a0_d3,a0_d4,a0_d5,a0_d6,a0_d7;
    reg [7:0] b0_d0,b0_d1,b0_d2,b0_d3,b0_d4,b0_d5,b0_d6,b0_d7;
    reg [7:0] a1_d0,a1_d1,a1_d2,a1_d3,a1_d4,a1_d5,a1_d6,a1_d7;
    reg [7:0] b1_d0,b1_d1,b1_d2,b1_d3,b1_d4,b1_d5,b1_d6,b1_d7;
    reg          wb_valid_r;
    reg [DW-1:0] wb_a0_r, wb_b0_r, wb_a1_r, wb_b1_r;
    reg [7:0]    wb_a0a_r, wb_b0a_r, wb_a1a_r, wb_b1a_r;

    reg [DW-1:0] coeff_out_r;
    reg          valid_o_r;
    reg          scale_valid_i;
    reg [DW-1:0] scale_a_i;

    wire [DW-1:0] zeta;
    zeta_rom u_z (.addr(zeta_idx), .zeta(zeta));

    wire [8:0] group_end        = {1'b0, start} + {1'b0, layer_len};
    wire [8:0] next_group_start = {1'b0, start} + {1'b0, layer_len} + {1'b0, layer_len};
    // pair (j, j+1) is the last in this group when j+2 reaches group_end
    wire       issue_group_last = ({1'b0, j} + 9'd2 >= group_end);
    wire       issue_layer_last = issue_group_last && (next_group_start >= 9'd256);
    wire [7:0] next_layer_len   = mode_r ? (layer_len << 1) : (layer_len >> 1);
    wire       issue_fire       = (state == S_CMP_RUN) && !layer_issue_done;

    // lane0 = position j, lane1 = position j+1 (same group => same zeta)
    wire [DW-1:0] issue_a0 = cr[j];
    wire [DW-1:0] issue_b0 = cr[j + layer_len];
    wire [DW-1:0] issue_a1 = cr[j + 8'd1];
    wire [DW-1:0] issue_b1 = cr[j + 8'd1 + layer_len];

    wire [DW-1:0] bf0_a_out, bf0_b_out, bf1_a_out, bf1_b_out;
    wire          bf0_valid, bf1_valid;
    butterfly_unit_pipe u_bf0 (
        .clk(clk), .rst_n(rst_n),
        .valid_i(issue_fire),
        .a(issue_a0), .b(issue_b0),
        .zeta(zeta), .mode(mode_r),
        .a_out(bf0_a_out), .b_out(bf0_b_out), .valid_o(bf0_valid)
    );
    butterfly_unit_pipe u_bf1 (
        .clk(clk), .rst_n(rst_n),
        .valid_i(issue_fire),
        .a(issue_a1), .b(issue_b1),
        .zeta(zeta), .mode(mode_r),
        .a_out(bf1_a_out), .b_out(bf1_b_out), .valid_o(bf1_valid)
    );

    wire [DW-1:0] scale_product;
    wire          scale_valid_o;
    barrett_mul_pipe u_scl (
        .clk(clk), .rst_n(rst_n),
        .valid_i(scale_valid_i),
        .a(scale_a_i), .b(12'd3303),
        .product(scale_product), .valid_o(scale_valid_o)
    );

    assign ready_o   = (state == S_IDLE) || (state == S_LOAD);
    assign coeff_out = coeff_out_r;
    assign valid_o   = valid_o_r;
    assign done_o    = (state == S_DONE);

    always @* begin
        next_state = state;
        case (state)
            S_IDLE:      if (valid_i) next_state = S_LOAD;
            S_LOAD:      if (load_cnt >= 8'd255 && valid_i) next_state = S_CMP_RUN;
            S_CMP_RUN:   if (layer_issue_done) next_state = S_CMP_DRAIN;
            S_CMP_DRAIN: if (bf_inflight == 6'd0)
                              next_state = (layer + 3'd1 >= LAYERS) ? S_OUT_PREP : S_CMP_RUN;
            S_OUT_PREP:  next_state = mode_r ? S_OUT_SCALE : S_OUTPUT;
            S_OUTPUT:    if (valid_o_r && ready_i && out_cnt >= 8'd255) next_state = S_DONE;
            S_OUT_SCALE: if (scale_emit_cnt >= 9'd256) next_state = S_DONE;
            S_DONE:      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;
            out_cnt <= 8'd0;
            scale_issue_cnt <= 9'd0;
            scale_emit_cnt <= 9'd0;
            j <= 8'd0;
            start <= 8'd0;
            layer_len <= 8'd0;
            zeta_idx <= 7'd0;
            layer <= 3'd0;
            bf_done <= 1'b0;
            mode_r <= 1'b0;
            layer_issue_done <= 1'b0;
            bf_inflight <= 6'd0;
            a0_d0<=8'd0;a0_d1<=8'd0;a0_d2<=8'd0;a0_d3<=8'd0;a0_d4<=8'd0;a0_d5<=8'd0;a0_d6<=8'd0;a0_d7<=8'd0;
            b0_d0<=8'd0;b0_d1<=8'd0;b0_d2<=8'd0;b0_d3<=8'd0;b0_d4<=8'd0;b0_d5<=8'd0;b0_d6<=8'd0;b0_d7<=8'd0;
            a1_d0<=8'd0;a1_d1<=8'd0;a1_d2<=8'd0;a1_d3<=8'd0;a1_d4<=8'd0;a1_d5<=8'd0;a1_d6<=8'd0;a1_d7<=8'd0;
            b1_d0<=8'd0;b1_d1<=8'd0;b1_d2<=8'd0;b1_d3<=8'd0;b1_d4<=8'd0;b1_d5<=8'd0;b1_d6<=8'd0;b1_d7<=8'd0;
            wb_valid_r <= 1'b0;
            wb_a0_r <= 12'd0; wb_b0_r <= 12'd0; wb_a1_r <= 12'd0; wb_b1_r <= 12'd0;
            wb_a0a_r <= 8'd0; wb_b0a_r <= 8'd0; wb_a1a_r <= 8'd0; wb_b1a_r <= 8'd0;
            coeff_out_r <= 12'd0;
            valid_o_r <= 1'b0;
            scale_valid_i <= 1'b0;
            scale_a_i <= 12'd0;
            for (ci = 0; ci < N; ci = ci + 1) cr[ci] = 12'd0;
        end else begin
            state <= next_state;
            scale_valid_i <= 1'b0;

            // Writeback address delay lines (8-deep = butterfly latency).
            a0_d0<=j;                a0_d1<=a0_d0;a0_d2<=a0_d1;a0_d3<=a0_d2;a0_d4<=a0_d3;a0_d5<=a0_d4;a0_d6<=a0_d5;a0_d7<=a0_d6;
            b0_d0<=j+layer_len;      b0_d1<=b0_d0;b0_d2<=b0_d1;b0_d3<=b0_d2;b0_d4<=b0_d3;b0_d5<=b0_d4;b0_d6<=b0_d5;b0_d7<=b0_d6;
            a1_d0<=j+8'd1;           a1_d1<=a1_d0;a1_d2<=a1_d1;a1_d3<=a1_d2;a1_d4<=a1_d3;a1_d5<=a1_d4;a1_d6<=a1_d5;a1_d7<=a1_d6;
            b1_d0<=j+8'd1+layer_len; b1_d1<=b1_d0;b1_d2<=b1_d1;b1_d3<=b1_d2;b1_d4<=b1_d3;b1_d5<=b1_d4;b1_d6<=b1_d5;b1_d7<=b1_d6;

            // In-flight pair accounting (mirrors single-butterfly version).
            if (issue_fire && !wb_valid_r)
                bf_inflight <= bf_inflight + 6'd1;
            else if (!issue_fire && wb_valid_r)
                bf_inflight <= bf_inflight - 6'd1;

            if (state != S_OUTPUT && state != S_OUT_SCALE)
                valid_o_r <= 1'b0;

            if (state == S_IDLE && valid_i) begin
                cr[0] <= coeff_in;
                load_cnt <= 8'd1;
                out_cnt <= 8'd0;
                scale_issue_cnt <= 9'd0;
                scale_emit_cnt <= 9'd0;
                j <= 8'd0;
                start <= 8'd0;
                layer <= 3'd0;
                bf_done <= 1'b0;
                layer_issue_done <= 1'b0;
                bf_inflight <= 6'd0;
                mode_r <= mode;
                if (!mode) begin
                    layer_len <= 8'd128;
                    zeta_idx  <= 7'd1;
                end else begin
                    layer_len <= 8'd2;
                    zeta_idx  <= 7'd127;
                end
            end

            if (state == S_LOAD && valid_i) begin
                cr[load_cnt] <= coeff_in;
                load_cnt <= load_cnt + 8'd1;
            end

            // ---- issue advance: step j by 2 (two butterflies/cycle) ----
            if (state == S_CMP_RUN && !layer_issue_done) begin
                if (issue_group_last) begin
                    if (!mode_r) zeta_idx <= zeta_idx + 7'd1;
                    else         zeta_idx <= zeta_idx - 7'd1;

                    if (issue_layer_last) begin
                        layer_issue_done <= 1'b1;
                    end else begin
                        start <= next_group_start[7:0];
                        j <= next_group_start[7:0];
                    end
                end else begin
                    j <= j + 8'd2;
                end
            end

            // ---- writeback: commit previous cycle's pair results ----
            if (wb_valid_r) begin
                cr[wb_a0a_r] <= wb_a0_r;
                cr[wb_b0a_r] <= wb_b0_r;
                cr[wb_a1a_r] <= wb_a1_r;
                cr[wb_b1a_r] <= wb_b1_r;
            end

            wb_valid_r <= bf0_valid;   // both lanes share latency/valid
            if (bf0_valid) begin
                wb_a0_r <= bf0_a_out; wb_b0_r <= bf0_b_out;
                wb_a1_r <= bf1_a_out; wb_b1_r <= bf1_b_out;
                wb_a0a_r <= a0_d7; wb_b0a_r <= b0_d7;
                wb_a1a_r <= a1_d7; wb_b1a_r <= b1_d7;
            end

            // ---- layer advance after drain ----
            if (state == S_CMP_DRAIN && bf_inflight == 6'd0) begin
                if (layer + 3'd1 >= LAYERS) begin
                    bf_done <= 1'b1;
                end else begin
                    layer <= layer + 3'd1;
                    layer_len <= next_layer_len;
                    start <= 8'd0;
                    j <= 8'd0;
                    layer_issue_done <= 1'b0;
                end
            end

            if (state == S_OUT_PREP) begin
                out_cnt <= 8'd0;
                scale_issue_cnt <= 9'd0;
                scale_emit_cnt <= 9'd0;
                if (!mode_r) begin
                    coeff_out_r <= cr[0];
                    valid_o_r <= 1'b1;
                end
            end

            if (state == S_OUTPUT && valid_o_r && ready_i) begin
                if (out_cnt < 8'd255) begin
                    out_cnt <= out_cnt + 8'd1;
                    coeff_out_r <= cr[out_cnt + 8'd1];
                    valid_o_r <= 1'b1;
                end else begin
                    out_cnt <= 8'd0;
                    valid_o_r <= 1'b0;
                end
            end

            if (state == S_OUT_SCALE) begin
                if (scale_issue_cnt < 9'd256) begin
                    scale_valid_i <= 1'b1;
                    scale_a_i <= cr[scale_issue_cnt[7:0]];
                    scale_issue_cnt <= scale_issue_cnt + 9'd1;
                end
                valid_o_r <= 1'b0;
                if (scale_valid_o) begin
                    coeff_out_r <= scale_product;
                    valid_o_r <= 1'b1;
                    scale_emit_cnt <= scale_emit_cnt + 9'd1;
                end
            end

            if (state == S_DONE) begin
                load_cnt <= 8'd0;
                out_cnt <= 8'd0;
                scale_issue_cnt <= 9'd0;
                scale_emit_cnt <= 9'd0;
                layer_issue_done <= 1'b0;
                bf_inflight <= 6'd0;
                wb_valid_r <= 1'b0;
                valid_o_r <= 1'b0;
            end
        end
    end
endmodule