`timescale 1ps / 1ps

module mem_write_B
#(
    parameter integer N2           = 4,
    parameter integer MATRIXSIZE_W = 16,
    parameter integer ADDR_W       = 12
)
(
    input  logic                    clk,
    input  logic                    rst,
    input  logic [MATRIXSIZE_W-1:0] M2,
    input  logic [MATRIXSIZE_W-1:0] M3dN2,
    input  logic                    valid_B,
    output logic  [ADDR_W-1:0]       wr_addr_B,
    output logic  [N2-1:0]           activate_B
);

/*
Each bank has M3dN2 columns
activate_B for the banks
wr_addr_B for the address within the banks
*/

logic [MATRIXSIZE_W-1:0] row;
logic [MATRIXSIZE_W-1:0] row_div;
logic [MATRIXSIZE_W-1:0] row_mod;
logic [MATRIXSIZE_W-1:0] col;

always_ff @(posedge clk) begin
    if(rst) begin
        wr_addr_B <= 0;
        activate_B <= 0;
        row <= 0;
        row_div <= 0;
        row_mod <= 0;
        col <= 0;
    end else begin
        if(valid_B) begin
            if (col < N2 - 1) begin
                col <= col + 1;
            end else begin
                col <= 0;
                if(row_mod == M3dN2 - 1) begin
                    row <= 0;
                    row_mod <= 0;
                    row_div <= row_div + 1;
                end else begin
                    row <= row + M2;
                    row_mod <= row_mod + 1;
                end
            end

        // wr_addr_B <= M2 * (row % M3dN2) + row_div;
        wr_addr_B <= row + row_div;
        activate_B <= 1 << col;
        end else begin
            wr_addr_B <= 0;
            activate_B <= 0;
        end
    end
end

endmodule