// ntt_core.v - NTT core with 4-way banked BRAM coefficient storage
//
// Coefficient storage moved from a 256x12 LUT register file to coeff_bram_4way
// (4 x 128x12 sd_bram, conflict-free banked). This removes the large multi-port
// LUT mux tree. Two pipelined butterflies (R1) still issue one pair/cycle.
//
// Because BRAM reads are synchronous (1-cycle latency), the compute pipeline
// gains one stage:
//   READ (present rd addrs from j) -> FEED butterflies (t+1, rd_data valid)
//   -> 8-stage Barrett butterfly -> WRITEBACK (t+9) into the write buffer.
// Read and writeback hit disjoint banks (butterfly index sets in a layer are
// disjoint), so each bank does 1 read + 1 write per cycle (sd_bram = 1R1W).
//
// LOAD/OUTPUT/OUT_SCALE stream through a single bank port with 1-cycle read
// latency (ready_i is tied high by all NTT consumers, so no output backpressure).

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;

    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 [8:0] out_rptr;        // output read address pointer (0..256)
    reg [8:0] out_cnt;         // output coeffs emitted (0..256)
    reg [8:0] scale_rptr;      // scale read address pointer
    reg [8:0] scale_issue_cnt; // coeffs fed to scale barrett
    reg [8:0] scale_emit_cnt;  // coeffs emitted from scale barrett
    reg [7:0] j, start, layer_len;
    reg [6:0] zeta_idx;
    reg [2:0] layer;
    reg       mode_r;
    reg       layer_issue_done;
    reg [5:0] bf_inflight;

    reg       rd_fire_r;       // issue_fire delayed 1 (BRAM read latency) -> feed
    // 9-deep writeback address delay lines (BRAM read + 8-stage butterfly).
    reg [7:0] a0d[0:8]; reg [7:0] b0d[0:8];
    reg [7:0] a1d[0:8]; reg [7:0] b1d[0:8];
    integer k;

    reg [DW-1:0] coeff_out_r;
    reg          valid_o_r;
    reg          scale_valid_i;
    reg [DW-1:0] scale_a_i;
    reg [DW-1:0] zeta_fed;     // zeta aligned to BRAM read data (1 cyc after present)

    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};
    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;

    // ---- coeff_bram_4way port muxing ----
    // Read addr 0 is shared by compute (j), output (out_rptr), scale (scale_rptr).
    wire [7:0] cmp_a0 = j;
    wire [7:0] cmp_b0 = j + layer_len;
    wire [7:0] cmp_a1 = j + 8'd1;
    wire [7:0] cmp_b1 = j + 8'd1 + layer_len;

    wire [7:0] rd_addr_0 = (state == S_OUT_PREP || state == S_OUTPUT) ? out_rptr[7:0]   :
                           (state == S_OUT_SCALE)                     ? scale_rptr[7:0] : cmp_a0;
    // Ports 1..3 carry the other three butterfly operands ONLY during compute;
    // otherwise they must mirror port 0 so the read crossbar's per-bank priority
    // arbitration doesn't route a stale compute address into port 0's bank.
    wire cmp_rd = (state == S_CMP_RUN) || (state == S_CMP_DRAIN);
    wire [7:0] rd_addr_1 = cmp_rd ? cmp_b0 : rd_addr_0;
    wire [7:0] rd_addr_2 = cmp_rd ? cmp_a1 : rd_addr_0;
    wire [7:0] rd_addr_3 = cmp_rd ? cmp_b1 : rd_addr_0;
    wire [11:0] rd_data_0, rd_data_1, rd_data_2, rd_data_3;

    // Butterflies fed from BRAM read data one cycle after the address (rd_fire_r).
    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(rd_fire_r),
        .a(rd_data_0), .b(rd_data_1),
        .zeta(zeta_fed), .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(rd_fire_r),
        .a(rd_data_2), .b(rd_data_3),
        .zeta(zeta_fed), .mode(mode_r),
        .a_out(bf1_a_out), .b_out(bf1_b_out), .valid_o(bf1_valid)
    );
    // zeta must align with the BRAM read data (1 cycle after presentation).

    // Scale multiplier (inverse NTT final x3303).
    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)
    );

    // ---- write ports ----
    wire loading_we = (state == S_IDLE && valid_i) || (state == S_LOAD && valid_i);
    wire [7:0] load_addr = (state == S_IDLE) ? 8'd0 : load_cnt;

    wire wr_en_0 = loading_we | bf0_valid;
    wire wr_en_1 = bf0_valid & ~loading_we;
    wire wr_en_2 = bf0_valid & ~loading_we;
    wire wr_en_3 = bf0_valid & ~loading_we;
    wire [7:0]  wr_addr_0 = loading_we ? load_addr : a0d[8];
    wire [7:0]  wr_addr_1 = b0d[8];
    wire [7:0]  wr_addr_2 = a1d[8];
    wire [7:0]  wr_addr_3 = b1d[8];
    wire [11:0] wr_data_0 = loading_we ? coeff_in : bf0_a_out;
    wire [11:0] wr_data_1 = bf0_b_out;
    wire [11:0] wr_data_2 = bf1_a_out;
    wire [11:0] wr_data_3 = bf1_b_out;

    coeff_bram_4way u_cr (
        .clk(clk), .rst_n(rst_n),
        .rd_addr_0(rd_addr_0), .rd_addr_1(rd_addr_1),
        .rd_addr_2(rd_addr_2), .rd_addr_3(rd_addr_3),
        .rd_data_0(rd_data_0), .rd_data_1(rd_data_1),
        .rd_data_2(rd_data_2), .rd_data_3(rd_data_3),
        .wr_en_0(wr_en_0), .wr_en_1(wr_en_1), .wr_en_2(wr_en_2), .wr_en_3(wr_en_3),
        .wr_addr_0(wr_addr_0), .wr_addr_1(wr_addr_1),
        .wr_addr_2(wr_addr_2), .wr_addr_3(wr_addr_3),
        .wr_data_0(wr_data_0), .wr_data_1(wr_data_1),
        .wr_data_2(wr_data_2), .wr_data_3(wr_data_3)
    );

    assign ready_o   = (state == S_IDLE) || (state == S_LOAD);
    // During forward OUTPUT, drive coeff_out straight from the BRAM read (1-cycle
    // latency, matching the old combinational-read + 1 output reg). Scale path
    // uses the registered barrett result.
    assign coeff_out = (state == S_OUTPUT) ? rd_data_0 : 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 (out_cnt >= 9'd256) 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_rptr <= 9'd0; out_cnt <= 9'd0;
            scale_rptr <= 9'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; zeta_fed <= 12'd0;
            layer <= 3'd0; mode_r <= 1'b0;
            layer_issue_done <= 1'b0; bf_inflight <= 6'd0;
            rd_fire_r <= 1'b0;
            for (k = 0; k <= 8; k = k + 1) begin
                a0d[k] <= 8'd0; b0d[k] <= 8'd0; a1d[k] <= 8'd0; b1d[k] <= 8'd0;
            end
            coeff_out_r <= 12'd0; valid_o_r <= 1'b0;
            scale_valid_i <= 1'b0; scale_a_i <= 12'd0;
        end else begin
            state <= next_state;
            scale_valid_i <= 1'b0;

            // read->feed pipeline register (BRAM latency) + zeta alignment
            rd_fire_r <= issue_fire;
            zeta_fed  <= zeta;

            // writeback address delay lines: sample j-derived addrs at read cycle
            a0d[0] <= cmp_a0; b0d[0] <= cmp_b0; a1d[0] <= cmp_a1; b1d[0] <= cmp_b1;
            for (k = 1; k <= 8; k = k + 1) begin
                a0d[k] <= a0d[k-1]; b0d[k] <= b0d[k-1];
                a1d[k] <= a1d[k-1]; b1d[k] <= b1d[k-1];
            end

            // in-flight pair accounting: +1 on issue, -1 on writeback
            if (issue_fire && !bf0_valid)      bf_inflight <= bf_inflight + 6'd1;
            else if (!issue_fire && bf0_valid) 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
                load_cnt <= 8'd1;
                out_rptr <= 9'd0; out_cnt <= 9'd0;
                scale_rptr <= 9'd0; scale_issue_cnt <= 9'd0; scale_emit_cnt <= 9'd0;
                j <= 8'd0; start <= 8'd0; layer <= 3'd0;
                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) load_cnt <= load_cnt + 8'd1;

            // ---- issue advance (read pointer j, step 2) ----
            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) layer_issue_done <= 1'b1;
                    else begin start <= next_group_start[7:0]; j <= next_group_start[7:0]; end
                end else begin
                    j <= j + 8'd2;
                end
            end

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

            // ---- OUT_PREP: reset output pointers ----
            if (state == S_OUT_PREP) begin
                out_rptr <= 9'd0; out_cnt <= 9'd0;
                scale_rptr <= 9'd1; scale_issue_cnt <= 9'd0; scale_emit_cnt <= 9'd0;
                // (scale first addr 0 is presented via out_rptr in this PREP cycle)
            end

            // ---- OUTPUT: stream coeffs. coeff_out is combinational rd_data_0;
            // valid_o is a registered "presented a valid addr last cycle". Present
            // addresses out_rptr = 0..255 (256), giving 256 valid outputs. ----
            if (state == S_OUTPUT) begin
                if (out_rptr < 9'd256) begin
                    out_rptr  <= out_rptr + 9'd1;   // present rd_addr_0 = out_rptr
                    valid_o_r <= 1'b1;              // data valid next cycle
                end else begin
                    valid_o_r <= 1'b0;
                end
                if (valid_o_r) out_cnt <= out_cnt + 9'd1;   // count emitted
            end

            // ---- OUT_SCALE: read coeff, x3303 via barrett, emit ----
            if (state == S_OUT_SCALE) begin
                if (scale_rptr < 9'd256) scale_rptr <= scale_rptr + 9'd1;
                // feed barrett with the coeff that arrived this cycle
                if (scale_issue_cnt < 9'd256) begin
                    scale_valid_i <= 1'b1;
                    scale_a_i <= rd_data_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_rptr <= 9'd0; out_cnt <= 9'd0;
                scale_rptr <= 9'd0; scale_issue_cnt <= 9'd0; scale_emit_cnt <= 9'd0;
                layer_issue_done <= 1'b0; bf_inflight <= 6'd0;
                valid_o_r <= 1'b0;
            end
        end
    end
endmodule