// ntt_core.v - NTT core with individual coefficient registers
//
// Uses 256 individual 12-bit registers and generate-based muxing
// to avoid any part-select simulation issues.
// 3-cycle butterfly: SetAddr -> Read -> Compute+Write

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;

    // Individual coefficient registers
    reg [DW-1:0] cr [0:N-1];
    integer ci;

    // State machine
    localparam S_IDLE=3'd0, S_LOAD=3'd1, S_CMP_A=3'd2, S_CMP_B=3'd3,
               S_CMP_C=3'd4, S_OUTPUT=3'd5, S_DONE=3'd6;
    reg [2:0] state, next_state;

    reg [7:0] load_cnt, out_cnt;
    reg [7:0] j, start, layer_len;
    reg [6:0] zeta_idx;
    reg [2:0] layer;
    reg bf_done;

    // Pipeline registers
    reg [DW-1:0] r_a, r_b;
    reg [7:0] r_wa, r_wb;

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

    // Butterfly
    wire [DW-1:0] bf_a_out, bf_b_out;
    butterfly_unit u_bf (
        .a(r_a), .b(r_b), .zeta(zeta), .mode(mode),
        .a_out(bf_a_out), .b_out(bf_b_out));

    // Output scaling
    wire [DW-1:0] coeff_scaled;
    barrett_mul u_scl (.a(cr[out_cnt]), .b(12'd3303), .product(coeff_scaled));
    assign coeff_out = mode ? coeff_scaled : cr[out_cnt];


    assign ready_o = (state == S_IDLE) || (state == S_LOAD);
    assign valid_o = (state == S_OUTPUT);
    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 >= 255 && valid_i) next_state = S_CMP_A;
            S_CMP_A: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_B;
            S_CMP_B: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_C;
            S_CMP_C: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_A;
            S_OUTPUT:if (out_cnt >= 255 && ready_i) 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<=0; out_cnt<=0; j<=0; start<=0; layer_len<=0;
            zeta_idx<=0; layer<=0; bf_done<=0; r_a<=0; r_b<=0; r_wa<=0; r_wb<=0;
            for (ci=0; ci<N; ci=ci+1) cr[ci] <= 0;
        end else begin
            state <= next_state;

            if (state == S_IDLE && valid_i) begin
                cr[0] <= coeff_in;
                load_cnt<=1; out_cnt<=0; j<=0; start<=0; layer<=0; bf_done<=0;
                if (!mode) begin layer_len<=128; zeta_idx<=1; end
                else       begin layer_len<=2;   zeta_idx<=127; end
            end

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

            // S_CMP_A: set read addresses (j, j+len)
            if (state == S_CMP_A) begin
                r_wa <= j;
                r_wb <= j + layer_len;
            end

            // S_CMP_B: capture read data
            if (state == S_CMP_B) begin
                r_a <= cr[j];
                r_b <= cr[j + layer_len];
            end

            // S_CMP_C: write butterfly results, advance counters
            if (state == S_CMP_C) begin
                cr[r_wa] <= bf_a_out;
                cr[r_wb] <= bf_b_out;

                j <= j + 8'd1;
                if (j + 8'd1 >= start + layer_len) begin
                    if (!mode) zeta_idx <= zeta_idx + 7'd1;
                    else       zeta_idx <= zeta_idx - 7'd1;

                    if ({1'b0,start} + {1'b0,layer_len} + {1'b0,layer_len} >= 256) begin
                        layer <= layer + 3'd1;
                        layer_len <= mode ? (layer_len<<1) : (layer_len>>1);
                        start <= 0; j <= 0;
                        if (layer + 3'd1 >= LAYERS) bf_done <= 1'b1;
                    end else begin
                        start <= start + layer_len + layer_len;
                        j <= start + layer_len + layer_len;
                    end
                end
            end

            if (state == S_OUTPUT && ready_i)
                out_cnt <= (out_cnt>=255) ? 0 : (out_cnt+8'd1);
        end
    end

endmodule