module matrix_memory_flexible #(
    parameter MEM_ROWS = 4,
    parameter MEM_COLS = 4,
    parameter INIT_FILE = "",
    parameter DATA_WIDTH = 16,
    parameter COLS_USED = 4,
    parameter OUTPUT_FILE = "memory_dump.hex"  // New parameter
)(
    input clk,

    // Control signals
    input write_enable,
    input read_enable,
    input read_full_row,  // 1 = read full row, 0 = read single element
    input write_full_row,
    input [$clog2(MEM_ROWS)-1:0] row,
    input [$clog2(MEM_COLS)-1:0] col,
    input read_full_row_or_col, //0 to read row wise & 1 to read col wise 
    input [$clog2(MEM_COLS)-1:0] no_cols_used, // in row wise reading or writing 
    input [$clog2(MEM_ROWS)-1:0] no_rows_used, // in col wise reading or writing 

    // Data ports
    input  [DATA_WIDTH-1:0] data_in,
    input   [DATA_WIDTH*((MEM_ROWS>MEM_COLS)?MEM_ROWS-1:MEM_COLS-1):0] full_row_in, //MEM_COLS_USED
    input write_back_to_file_enable, //writing enable the file back into the memory  

    output reg [DATA_WIDTH-1:0] data_out,
    output reg [DATA_WIDTH*((MEM_ROWS>MEM_COLS)?MEM_ROWS-1:MEM_COLS-1):0] full_row_out, //MEM_COLS_USED


    // Output valid signal
    output reg valid,

    // Writeback control
    output reg done_writing_to_file ////writing enable the file back into the memory
);

    // === 2D Memory Declaration ===
    reg [DATA_WIDTH-1:0] mem [0:MEM_ROWS-1][0:MEM_COLS-1];

    integer i, j;
    integer f, r, c;


    // === Optional Memory Initialization ===
    initial begin
        if (INIT_FILE != "") begin
            $display("Loading memory from: %s", INIT_FILE);
            $readmemh(INIT_FILE, mem);
        end else begin
            for (i = 0; i < MEM_ROWS; i = i + 1)
                for (j = 0; j < MEM_COLS; j = j + 1)
                    mem[i][j] = {DATA_WIDTH{1'b0}};
        end
        full_row_out <= 0;
        data_out <= 0;
        valid <= 0;
        done_writing_to_file <= 0;
    end

    // === Read/Write Logic ===
    always @(posedge clk) begin
        if (valid == 1)
            valid <= 0;
        
        if (write_full_row) begin////0 to read row wise & 1 to read col wise
            for (i = 0; i <= ((read_full_row_or_col == 0)?no_cols_used:no_rows_used); i = i + 1) begin//<
                if(read_full_row_or_col == 0) begin
                    mem[row][i + col] <= full_row_in[i*DATA_WIDTH +: DATA_WIDTH];
                    // $display("(i = %0h/%0h)Full Row Output(row wise ):%0h", i, no_cols_used, full_row_in);

                end else begin
                    mem[i + row][col] <= full_row_in[i*DATA_WIDTH +: DATA_WIDTH];
                    // $display("(i = %0h/%0h)Full Row Output(col wise ):%0h", i, no_rows_used, full_row_in);
                end
            end
        end else if (write_enable) begin
            mem[row][col] <= data_in;
        end else if (read_full_row && (valid == 0)) begin////0 to read row wise & 1 to read col wise
            for (i = 0; i <= ((read_full_row_or_col == 0)?no_cols_used:no_rows_used); i = i + 1) begin//<
                if(read_full_row_or_col == 0)begin
                    full_row_out[i*DATA_WIDTH +: DATA_WIDTH] <= mem[row][i + col];
                    // $display("(i = %0h/%0h)Full Row Output(row wise ):%0h; mem =%0h", i, no_cols_used, full_row_out, mem[row][i + col]);
                end else begin
                    full_row_out[i*DATA_WIDTH +: DATA_WIDTH] <= mem[i + row][col];
                    // $display("(i = %0h/%0h)Full Row Output(col wise ):%0h; mem =%0h", i, no_rows_used, full_row_out, mem[row][i + col]);
                end
            end
            $display("Full Row Output: %h", full_row_out);
            valid <= 1;
        end else if (read_enable) begin
            data_out <= mem[row][col];
            valid <= 1;
        end
        else if (valid == 0)
                full_row_out = {(DATA_WIDTH*((MEM_ROWS>MEM_COLS)?MEM_ROWS:MEM_COLS)){1'b0}};  // // not working to reset the full_row_out before next writing 

        // === Writeback logic (for simulation only) ===
        if (write_back_to_file_enable && !done_writing_to_file) begin
            f = $fopen(OUTPUT_FILE, "w");
            if (f) begin
                for (r = 0; r < MEM_ROWS; r = r + 1) begin
                    for (c = 0; c < MEM_COLS; c = c + 1) begin
                        $fwrite(f, "%04x\n", mem[r][c]);
                    end
                    // $fwrite(f, "\n");
                end
                $fclose(f);
                $display("Memory written to file: %s", OUTPUT_FILE);
            end else begin
                $display("Error opening file: %s", OUTPUT_FILE);
            end
            done_writing_to_file <= 1;
        end else if (!write_back_to_file_enable) begin
            done_writing_to_file <= 0;  // Allow re-triggering
        end
    end

endmodule