// tb_ntt_core_xsim.v - Standard Verilog testbench for ntt_core targeting Vivado xsim
//
// Reads test vectors from a hex file using $readmemh.
// Each line is a single hex number encoding both mode and 256 input coefficients:
//   - bit [3075]: mode (0 = forward NTT, 1 = inverse NTT)
//   - bits [3074:3]: 256 x 12-bit input coefficients
//   - bits [2:0]: unused (padding)
//   - Total: 769 hex chars per line, NO spaces
//
// Drives ntt_core, streams in 256 coefficients via valid/ready handshake,
// waits for valid_o, streams out 256 transformed coefficients,
// and writes "RESULT: MODE INDEX COEFF_HEX" to the output file.
//
// ntt_core FSM: IDLE → LOAD → CMP_A/B/C (7 stages) → OUTPUT → DONE → IDLE
// valid_o is asserted during OUTPUT state (before done_o).
// done_o pulses for 1 cycle AFTER all outputs have been read.
//
// Parameters:
//   VECTOR_FILE - path to input hex file (default: "vectors/ntt_core_input.hex")
//   RESULT_FILE - path to output file    (default: "vectors/ntt_core_result.hex")
//
// Usage:
//   xvlog -sv ntt_core.v butterfly_unit.v barrett_mul.v zeta_rom.v tb_ntt_core_xsim.v
//   xelab tb_ntt_core_xsim -s tb_ntt_core_xsim
//   xsim tb_ntt_core_xsim -R

`timescale 1ns / 1ps

module tb_ntt_core_xsim;

    // ================================================================
    // Parameters
    // ================================================================
    parameter VECTOR_FILE   = "sync_rtl/ntt/TB/vectors/ntt_core_input.hex";
    parameter RESULT_FILE   = "sync_rtl/ntt/TB/vectors/ntt_core_result.hex";
    parameter MAX_VECTORS   = 64;
    parameter N_COEFFS      = 256;
    parameter TIMEOUT_CYCLES= 100000;

    // ================================================================
    // DUT signals
    // ================================================================
    reg              clk;
    reg              rst_n;
    reg  [11:0]      coeff_in;
    reg              valid_i;
    wire             ready_o;
    reg              mode;
    wire [11:0]      coeff_out;
    wire             valid_o;
    reg              ready_i;
    wire             done_o;

    // ================================================================
    // DUT instantiation
    // ================================================================
    ntt_core u_dut (
        .clk       (clk),
        .rst_n     (rst_n),
        .coeff_in  (coeff_in),
        .valid_i   (valid_i),
        .ready_o   (ready_o),
        .mode      (mode),
        .coeff_out (coeff_out),
        .valid_o   (valid_o),
        .ready_i   (ready_i),
        .done_o    (done_o)
    );

    // ================================================================
    // Clock generation: 100 MHz (10 ns period)
    // ================================================================
    initial clk = 1'b0;
    always #5 clk = ~clk;

    // ================================================================
    // Vector memory (loaded by $readmemh)
    // 3076 bits per word:
    //   bit [3075]:        mode (0=FWD NTT, 1=INV NTT)
    //   bits [3074:3]:     256 x 12-bit coefficients (coeff[0] at MSB)
    //   bits [2:0]:        unused padding
    // ================================================================
    reg [3075:0] vector_mem [0:MAX_VECTORS-1];
    integer      vec_count;
    integer      idx;
    integer      ci;
    integer      cycle_count;
    integer      result_fd;

    // Test result tracking
    integer      pass_count;
    integer      fail_count;

    // Temporary storage for coefficients
    reg [11:0]   input_coeffs  [0:N_COEFFS-1];
    reg [11:0]   output_coeffs [0:N_COEFFS-1];

    // ================================================================
    // Hex-to-ASCII conversion helper
    // ================================================================
    function [7:0] nibble_to_ascii;
        input [3:0] nibble;
        begin
            if (nibble < 4'd10)
                nibble_to_ascii = 8'h30 + {4'd0, nibble};  // '0'-'9'
            else
                nibble_to_ascii = 8'h41 + ({4'd0, nibble} - 4'd10); // 'A'-'F'
        end
    endfunction

    // ================================================================
    // Main test sequence
    // ================================================================
    initial begin
        // Count loaded vectors
        vec_count = 0;

        // Load vectors from hex file
        $readmemh(VECTOR_FILE, vector_mem);

        // Count entries (XSim leaves unloaded entries as 3076'hX)
        begin
            integer found_end;
            found_end = 0;
            for (idx = 0; idx < MAX_VECTORS; idx = idx + 1) begin
                if (!found_end && (vector_mem[idx] === 3076'hx ||
                                   vector_mem[idx] === 3076'hz))
                    found_end = 1;
                else if (!found_end)
                    vec_count = vec_count + 1;
            end
        end

        if (vec_count == 0) begin
            $display("ERROR: No vectors loaded from %s", VECTOR_FILE);
            $display("       Check that the file exists and format is correct.");
            $display("       Each line: 769 hex chars = {1 hex digit mode, 768 hex chars coeffs}");
            $finish;
        end

        $display("INFO: Loaded %0d test vectors from %s", vec_count, VECTOR_FILE);

        // Open result file
        result_fd = $fopen(RESULT_FILE, "w");
        if (result_fd == 0) begin
            $display("ERROR: Cannot open result file: %s", RESULT_FILE);
            $finish;
        end

        // Initialize DUT inputs
        mode     <= 1'b0;
        coeff_in <= 12'd0;
        valid_i  <= 1'b0;
        ready_i  <= 1'b1;  // always ready to accept output

        // Reset sequence: rst_n low for 3 cycles, then high
        rst_n <= 1'b0;
        repeat (3) @(posedge clk);
        rst_n <= 1'b1;
        @(posedge clk);

        pass_count = 0;
        fail_count = 0;

        // ============================================================
        // Process each vector
        // ============================================================
        for (idx = 0; idx < vec_count; idx = idx + 1) begin
            begin
                reg         vec_mode;
                integer     out_idx;
                integer     got_done;

                // Extract mode and coefficients from vector memory
                vec_mode = vector_mem[idx][3075];

                // Extract 256 input coefficients
                // coeff[0] at bits [3074:3063], coeff[255] at bits [11:0]
                // NOTE: bits[2:0] are unused padding below coeff[255]
                for (ci = 0; ci < N_COEFFS; ci = ci + 1) begin
                    input_coeffs[ci] = vector_mem[idx][(N_COEFFS-1-ci)*12 +: 12];
                end

                $display("INFO: Vector %0d - mode=%0d (0=FWD, 1=INV)", idx, vec_mode);

                // Set mode
                mode <= vec_mode;

                // ----------------------------------------------------
                // LOAD phase: stream in 256 coefficients via valid/ready
                // ----------------------------------------------------
                for (ci = 0; ci < N_COEFFS; ci = ci + 1) begin
                    // Wait for ready_o before driving next input
                    cycle_count = 0;
                    while (!ready_o && cycle_count < TIMEOUT_CYCLES) begin
                        @(posedge clk);
                        cycle_count = cycle_count + 1;
                    end
                    if (cycle_count >= TIMEOUT_CYCLES) begin
                        $display("ERROR: Timeout waiting for ready_o (input ci=%0d) on vector %0d",
                                 ci, idx);
                        fail_count = fail_count + 1;
                        ci = N_COEFFS;  // break out of input loop
                    end else begin
                        coeff_in <= input_coeffs[ci];
                        valid_i  <= 1'b1;
                        @(posedge clk);
                        valid_i  <= 1'b0;
                    end
                end

                // ----------------------------------------------------
                // OUTPUT phase: wait for valid_o, then stream out results
                // ntt_core asserts valid_o during S_OUTPUT state.
                // After all 256 outputs are consumed, done_o pulses.
                // ----------------------------------------------------
                got_done = 0;
                out_idx = 0;

                while (out_idx < N_COEFFS) begin
                    // Wait for valid_o or done_o
                    cycle_count = 0;
                    while (!valid_o && !done_o && cycle_count < TIMEOUT_CYCLES) begin
                        @(posedge clk);
                        cycle_count = cycle_count + 1;
                    end

                    if (cycle_count >= TIMEOUT_CYCLES) begin
                        $display("ERROR: Timeout waiting for output (out_idx=%0d) on vector %0d",
                                 out_idx, idx);
                        fail_count = fail_count + 1;
                        out_idx = N_COEFFS;  // break out
                    end else if (done_o) begin
                        // done_o before all outputs? This shouldn't happen
                        // unless outputs were consumed in previous test somehow.
                        // Treat as error.
                        $display("WARN: done_o asserted before output %0d on vector %0d",
                                 out_idx, idx);
                        got_done = 1;
                        @(posedge clk);  // consume done_o cycle
                    end else begin
                        // valid_o is high: capture output
                        output_coeffs[out_idx] = coeff_out;
                        out_idx = out_idx + 1;
                        @(posedge clk);  // ready_i=1, consume this output
                    end
                end

                // After all 256 outputs captured, check if done_o was seen
                // or wait briefly for it
                if (!got_done) begin
                    cycle_count = 0;
                    while (!done_o && cycle_count < 100) begin
                        @(posedge clk);
                        cycle_count = cycle_count + 1;
                    end
                    if (cycle_count >= 100) begin
                        $display("WARN: done_o not seen after all outputs on vector %0d", idx);
                    end else begin
                        @(posedge clk);  // consume done_o
                    end
                end

                pass_count = pass_count + 1;

                // Write result to output file
                // Format: "RESULT: MODE INDEX COEFF_HEX"
                $fwrite(result_fd, "RESULT: %0d %0d ", vec_mode, idx);
                // Write all 256 output coeffs as hex (3 hex chars each = 768 total)
                begin
                    integer oi;
                    reg [3:0] nib;
                    for (oi = 0; oi < N_COEFFS; oi = oi + 1) begin
                        // Write 3 hex chars per coefficient (12 bits)
                        nib = output_coeffs[oi][11:8];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                        nib = output_coeffs[oi][7:4];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                        nib = output_coeffs[oi][3:0];
                        $fwrite(result_fd, "%c", nibble_to_ascii(nib));
                    end
                end
                $fwrite(result_fd, "\n");

                // Extra cycle before next vector
                @(posedge clk);
            end  // inner begin block
        end

        // ============================================================
        // Summary
        // ============================================================
        $fclose(result_fd);

        $display("========================================");
        $display("TEST COMPLETE");
        $display("  Total vectors: %0d", vec_count);
        $display("  Passed:        %0d", pass_count);
        $display("  Failed:        %0d", fail_count);
        $display("  Results written to: %s", RESULT_FILE);
        $display("========================================");

        $finish;
    end

    // ================================================================
    // Timeout watchdog
    // ================================================================
    initial begin
        // TIMEOUT_CYCLES * 10ns per cycle * 10x margin
        #(TIMEOUT_CYCLES * 10 * 10);
        $display("FATAL: Global simulation timeout reached");
        $finish;
    end

endmodule
