AYAKA_Transformer/rtl/mask_A_Q_KV_tb.v

`timescale 1ns / 1ps

module tb_top_masking #(    
    parameter ROWS = 20,///16
    parameter COLS = 80,///12
    parameter ROWS_OUT = 20,///16
    parameter COLS_OUT = 80,///32
    parameter DATA_WIDTH = 16
)();

    // Clock and reset
    reg clk = 0;
    reg rst = 1;
    always #5 clk = ~clk;

    // Control
    reg enable_mask;

    // Addressing
    reg [$clog2(ROWS)-1:0] rows_start_add_reading, rows_start_add_writing, rows_size_reading, rows_size_writing;
    reg [$clog2(COLS)-1:0] cols_start_add_reading, cols_size_reading;
    reg [$clog2(COLS_OUT)-1:0] cols_start_add_writing, cols_size_writing;

    // Connections
    wire [$clog2(ROWS)-1:0] row_addr_out_read, row_addr_out_write;
    wire [$clog2(COLS)-1:0] col_addr_out_read;
    wire [$clog2(COLS_OUT)-1:0] col_addr_out_write;

    wire write_enable;
    wire read_enable_full_row;
    wire read_valid;
    wire valid_result;

    // Data
    wire [DATA_WIDTH - 1:0] data_output;

    reg read_full_row_or_col;
    wire [DATA_WIDTH*((ROWS>COLS)?ROWS:COLS)-1:0] full_row_output;

    reg  write_back_to_file_enable;
    wire done_writing_to_file;

    // === Instantiate memory ===
    top_module_mem #(
        .ROWS1(ROWS), .COLS1(COLS),
        .ROWS2(ROWS), .COLS2(COLS),
        .ROWS3(ROWS_OUT), .COLS3(COLS_OUT),
        .DATA_WIDTH(DATA_WIDTH)
    ) u_top_mem (
        .clk(clk),

        // Memory 1 (reading)
        .row_addr_1(row_addr_out_read),
        .col_addr_1(col_addr_out_read),
        .write_enable_1(),
        .read_enable_1(),
        .data_input_1(0), // Not writing to memory 1
        .data_output_1(),
        .valid_1(read_valid),
        .read_full_row_or_col1(read_full_row_or_col),
        .read_full_row_1(read_enable_full_row),
        .no_cols_used1(cols_size_reading),
        .no_rows_used1(rows_size_reading),
        .full_row_output_1(full_row_output),
        .full_row_input_1(),
        .write_full_row_1(),

        // Memory 2 (not used)
        .row_addr_2(), .col_addr_2(), .write_enable_2(), .read_enable_2(),
        .data_input_2(), .valid_2(), .read_full_row_or_col2(),
        .read_full_row_2(), .no_cols_used2(), .no_rows_used2(), .data_output_2(),
        .full_row_output_2(), .full_row_input_2(), .write_full_row_2(),

        // Memory 3 (writing)
        .row_addr_3(row_addr_out_write),
        .col_addr_3(col_addr_out_write),
        .write_enable_3(write_enable),
        .read_enable_3(),
        .data_input_3(data_output),
        .data_output_3(),
        .valid_3(),
        .read_full_row_or_col3(1'b0),
        .read_full_row_3(),
        .no_cols_used3(cols_size_writing),
        .no_rows_used3(rows_size_writing),
        .full_row_output_3(),
        .full_row_input_3(),
        .write_full_row_3(),

        .write_back_to_file_enable(write_back_to_file_enable),
        .done_writing_to_file(done_writing_to_file)
    );

    // === Instantiate masking unit ===
    top_module_mask #(
        .ROWS_READING(ROWS),
        .COLS_READING(COLS),
        .ROWS_WRITING(ROWS_OUT),
        .COLS_WRITING(COLS_OUT),
        .DATA_WIDTH(DATA_WIDTH)
    ) u_mask (
        .clk(clk),
        .rst(rst),
        .enable_mask(enable_mask),

        .read_valid(read_valid),

        .rows_start_add_reading(rows_start_add_reading),
        .cols_start_add_reading(cols_start_add_reading),
        .rows_start_add_writing(rows_start_add_writing),
        .cols_start_add_writing(cols_start_add_writing),

        .rows_size_reading(rows_size_reading),
        .cols_size_reading(cols_size_reading),
        .rows_size_writing(rows_size_writing),
        .cols_size_writing(cols_size_writing),

        .data_input_rows(full_row_output),

        .data_output(data_output),
        .write_enable(write_enable),

        .row_addr_out_read(row_addr_out_read),
        .col_addr_out_read(col_addr_out_read),
        .row_addr_out_write(row_addr_out_write),
        .col_addr_out_write(col_addr_out_write),

        .read_enable_full_row(read_enable_full_row),
        .valid_result(valid_result)
    );

    // === Stimulus ===
    initial begin
        $display("=== Masking Testbench Starting ===");

        read_full_row_or_col = 0;  // row-wise
        enable_mask = 0;
        rst = 1;
        write_back_to_file_enable = 0;

        #20;
        rst = 0;

        // Configure parameters
        rows_start_add_reading   <= 0;
        cols_start_add_reading   <= 0;
        rows_start_add_writing   <= 0;
        cols_start_add_writing   <= 0;

        rows_size_reading        <= 5'd19;///4'd15
        cols_size_reading        <= 5'd19;///4'd11
        rows_size_writing        <= 5'd19;//4'd15
        // cols_size_writing        <= ((COLS/DATA_WIDTH)+2);
        cols_size_writing        <= ((COLS/DATA_WIDTH));


        #10;
        enable_mask = 1;

        wait(valid_result);

        #50;
        write_back_to_file_enable = 1;

        wait(done_writing_to_file);

        $display("=== Masking Testbench Done ===");
        $finish;
    end

endmodule