import numpy as np
import torch
import paddle

# 更换不同的shape
shapes = [(30, 20), (100, 200), (5, 400), (20, 5), (5, 2), (100, 5)]

# 定义要测试的损失函数组合
loss_functions = {
    'U': {
        'torch': lambda u, s, v: torch.sum(torch.abs(u)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(u))
    },
    'S': {
        'torch': lambda u, s, v: torch.sum(torch.abs(s)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(s))
    },
    'V': {
        'torch': lambda u, s, v: torch.sum(torch.abs(v)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(v))
    },
    'U+V': {
        'torch': lambda u, s, v: torch.sum(torch.abs(u)) + torch.sum(torch.abs(v)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(u)) + paddle.sum(paddle.abs(v))
    },
    'U+S': {
        'torch': lambda u, s, v: torch.sum(torch.abs(u)) + torch.sum(torch.abs(s)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(u)) + paddle.sum(paddle.abs(s))
    },
    'S+V': {
        'torch': lambda u, s, v: torch.sum(torch.abs(s)) + torch.sum(torch.abs(v)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(s)) + paddle.sum(paddle.abs(v))
    },
    'U+S+V': {
        'torch': lambda u, s, v: torch.sum(torch.abs(u)) + torch.sum(torch.abs(s)) + torch.sum(torch.abs(v)),
        'paddle': lambda u, s, v: paddle.sum(paddle.abs(u)) + paddle.sum(paddle.abs(s)) + paddle.sum(paddle.abs(v))
    }
}

for shape in shapes:    
    # ===================== 1. 生成复数随机数据 =====================
    np.random.seed(42)  # 固定随机种子
    real_part = np.random.rand(*shape).astype(np.float64)
    imag_part = np.random.rand(*shape).astype(np.float64)
    np_complex_array = real_part + 1j * imag_part  # 构造复数矩阵

    for loss_name, funcs in loss_functions.items():
        # ===================== 2. PyTorch 计算 =====================
        # 转换为 PyTorch 复数张量
        torch_complex_tensor = torch.from_numpy(np_complex_array).to('cuda').requires_grad_(True)

        # SVD 前向计算 (复数输入)
        U_torch, S_torch, V_torch = torch.linalg.svd(torch_complex_tensor, full_matrices=False)

        # 计算损失
        loss_torch = funcs['torch'](U_torch, S_torch, V_torch)

        # 反向传播
        loss_torch.backward()
        torch_grad = torch_complex_tensor.grad.cpu().numpy()

        # ===================== 3. PaddlePaddle 计算 =====================
        # 转换为 Paddle 复数张量 (paddle.complex64)
        paddle_complex_tensor = paddle.to_tensor(np_complex_array, stop_gradient=False)

        # SVD 前向计算 (复数输入)
        U_paddle, S_paddle, V_paddle = paddle.linalg.svd(paddle_complex_tensor, full_matrices=False)

        # 计算损失
        loss_paddle = funcs['paddle'](U_paddle, S_paddle, V_paddle)

        # 反向传播
        loss_paddle.backward()
        paddle_grad = paddle_complex_tensor.grad.numpy()

        # 比较梯度
        are_close = np.allclose(torch_grad, paddle_grad, atol=1e-5, rtol=1e-5)
        print(f"梯度是否一致 (shape {shape}, loss {loss_name}):", are_close)
        
        # 如果不一致，打印最大差异
        if not are_close:
            max_diff = np.max(np.abs(torch_grad - paddle_grad))
            print(f"最大差异: {max_diff}")