import torch
from typing import List, Tuple


# def vec(K):
#     return K.T.flatten().reshape(-1, 1)

# def rebuild(K, r1, r2):
#     """
#     Implements the R(K) operation from the image.
#     K: input matrix (k x d)
#     r1: block height
#     r2: number of block columns
#     """
#     k, d = K.shape
#     num_block_rows = k // r1
#     num_block_cols = r2
#     bw = d // r2 # block width
    
#     blocks = []
#     # R(K) stacks vec(Ki,j) as columns. 
#     # The image shows column-major order through the blocks.
#     for j in range(num_block_cols):
#         for i in range(num_block_rows):
#             # Extract block Ki,j
#             Ki_j = K[i*r1:(i+1)*r1, j*bw:(j+1)*bw]
#             # Vectorize (column-major) and add to list
#             blocks.append(vec(Ki_j))
    
#     return torch.hstack(blocks)


def rebuild(K, r1, r2):
    k, d = K.shape
    Br = k // r1      # number of block rows
    bw = d // r2      # block width

    # Step 1: reshape to (Br, r1, r2, bw)
    K_view = K.view(Br, r1, r2, bw)

    # Step 2: we want to vectorize each (r1, bw) block in COLUMN-MAJOR order.
    # That is equivalent to transposing the block and flattening in row-major.
    # So we permute to (Br, r2, bw, r1) and then flatten last two dims.
    
    # But better: move r1 and bw to end, then transpose those two
    # Actually: to get column-major flatten of (r1, bw), we can do:
    #   block.transpose(-2, -1).contiguous().view(-1)
    # So let's transpose the last two dims of the block

    # Current: (Br, r1, r2, bw) → we want to treat (r1, bw) as block → transpose to (bw, r1)
    # So permute to (Br, r2, bw, r1)
    K_transposed_blocks = K_view.permute(0, 2, 3, 1)  # (Br, r2, bw, r1)

    # Now flatten the last two dims (bw, r1) → (bw * r1,) → this is column-major of original block
    vecs = K_transposed_blocks.reshape(Br, r2, bw * r1)  # (Br, r2, vec_len)

    # Now, we have vecs[i, j] = vectorized block (i,j)
    # But we want to output columns in order: j=0: i=0,1,...,Br-1; j=1: i=0,...
    # So we need to **transpose the first two dimensions** and then **flatten in row-major**
    
    # Transpose to (r2, Br, vec_len)
    vecs = vecs.permute(1, 0, 2)  # (r2, Br, vec_len)

    # Now flatten first two dims: (r2*Br, vec_len), then transpose to (vec_len, r2*Br)
    result = vecs.reshape(r2 * Br, -1).t()

    return result

# def rebuild(grad, block_size: [int, int]):
#     new_matrix_rows = []
#     if grad.dim() == 2:  # 只处理二维梯度（矩阵）
#         # 获取梯度矩阵的尺寸
#         rows, cols = grad.size()
#         # 遍历分块
#         for j in range(0, cols, block_size[1]):
#             for i in range(0, rows, block_size[0]):
#                 # 获取当前块
#                 block = grad[i:i + block_size[0], j:j + block_size[1]]
#                 # 如果块的大小不足，填充零
#                 if block.size(0) < block_size[0] or block.size(1) < block_size[1]:
#                     padding = (
#                         0, block_size[1] - block.size(1),  # 列填充
#                         0, block_size[0] - block.size(0)  # 行填充
#                     )
#                     block = torch.nn.functional.pad(block, padding, "constant", 0)
#                 # 向量化并添加到新矩阵的行中
#                 new_matrix_rows.append(block.T.flatten())

#         # 将所有行堆叠成一个新矩阵
#     if new_matrix_rows:  # 如果有数据
#         new_gad = torch.stack(new_matrix_rows)
#     else:
#         new_gad = torch.empty(0)  # 如果没有梯度数据，返回空矩阵

#     return new_gad

# if __name__ == "__main__":
#     # 定义一个简单的模型
#     class SimpleModel(torch.nn.Module):
#         def __init__(self):
#             super(SimpleModel, self).__init__()
#             self.fc1 = torch.nn.Linear(10, 20)
#             self.fc2 = torch.nn.Linear(20, 10)

#         def forward(self, x):
#             x = self.fc1(x)
#             x = self.fc2(x)
#             return x

#     # 初始化模型和损失函数
#     model = SimpleModel()
#     criterion = torch.nn.MSELoss()
#     optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

#     # 模拟输入数据
#     inputs = torch.randn(5, 10)  # batch_size=5, input_size=10
#     targets = torch.randn(5, 10)  # batch_size=5, output_size=10

#     # 前向传播
#     outputs = model(inputs)
#     loss = criterion(outputs, targets)

#     # 反向传播计算梯度
#     loss.backward()

#     # 调用函数将梯度分块并构造新矩阵
#     for param in model.parameters():
#         if param.grad is not None:  # 检查是否有梯度
#             grad = param.grad  # 获取梯度
#             print("旧矩阵的内容:\n", grad)
#             print("新矩阵的形状:", grad.shape)
#             block_size = (2, 2)  # 分块大小为 2x2
#             new_matrix = rebuild(grad, block_size)
#             print("新矩阵的形状:", new_matrix.shape)
#             print("新矩阵的内容:\n", new_matrix)