"""
UOT (Unbalanced Optimal Transport) 位姿估计 Demo
=================================================
UOTHead 使用 Sinkhorn 非平衡最优传输进行特征匹配和位姿估计。

流程:
1. Cosine Cost Matrix: C = 1 - cosine_sim(feat1, feat2)
2. Sinkhorn Unbalanced OT: 迭代求解运输计划 T
3. Point Projection: project_kpts = T @ kpts2 / sum(T)
4. Weighted SVD: 从匹配点对估计刚体变换 R|t

关键参数:
- epsilon: 熵正则化（控制运输计划的平滑度）
- gamma: 质量正则化（允许部分匹配）
- sinkhorn_iter: 5次迭代
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from uot import UOTHead, sinkhorn_unbalanced, compute_rigid_transform

OUTPUT_DIR = os.path.join(os.path.dirname(__file__), 'output')
os.makedirs(OUTPUT_DIR, exist_ok=True)


def visualize_cost_matrix():
    """可视化代价矩阵"""
    print('\n--- 代价矩阵 (Cost Matrix) ---')

    torch.manual_seed(42)
    # 模拟query和positive的150个关键点特征
    feat1 = torch.randn(2, 150, 128)  # query
    feat2 = torch.randn(2, 150, 128)  # positive

    # 让部分特征相似（模拟真实闭环场景）
    # 前100个特征点有对应关系
    feat2[:, :100] = feat1[:, :100] + 0.1 * torch.randn(2, 100, 128)

    # 计算cosine cost matrix
    feat1_norm = feat1 / (feat1.norm(dim=2, keepdim=True) + 1e-8)
    feat2_norm = feat2 / (feat2.norm(dim=2, keepdim=True) + 1e-8)
    C = 1.0 - torch.bmm(feat1_norm, feat2_norm.transpose(1, 2))

    C_np = C[0].numpy()

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    im0 = axes[0].imshow(C_np, cmap='YlOrRd')
    axes[0].set_title('Cost Matrix C = 1 - cos_sim')
    axes[0].set_xlabel('Positive Point j')
    axes[0].set_ylabel('Query Point i')
    plt.colorbar(im0, ax=axes[0])

    # 缩放看前30个点（有对应关系的）
    im1 = axes[1].imshow(C_np[:30, :30], cmap='YlOrRd')
    axes[1].set_title('Cost Matrix (前30×30)\n有模拟对应关系')
    axes[1].set_xlabel('Positive Point j')
    axes[1].set_ylabel('Query Point i')
    plt.colorbar(im1, ax=axes[1])

    # 对角线cost分布 vs 非对角线
    diag_cost = np.diag(C_np)
    off_diag = C_np[~np.eye(150, dtype=bool)]

    axes[2].hist(diag_cost, bins=30, alpha=0.6, label=f'对角线(匹配点)\nmean={diag_cost.mean():.3f}',
                 color='green')
    axes[2].hist(off_diag, bins=30, alpha=0.6, label=f'非对角线\nmean={off_diag.mean():.3f}',
                 color='gray')
    axes[2].set_title('Cost分布: 匹配 vs 非匹配')
    axes[2].set_xlabel('Cost')
    axes[2].legend(fontsize=8)

    plt.suptitle('UOT 代价矩阵分析', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'uot_cost_matrix.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'  [保存] {path}')


def visualize_sinkhorn():
    """可视化Sinkhorn迭代过程"""
    print('\n--- Sinkhorn 迭代过程 ---')

    torch.manual_seed(42)
    # 构造有明显对应关系的特征
    B, N, C = 1, 50, 128
    feat1 = torch.randn(B, N, C)
    feat1_norm = feat1 / (feat1.norm(dim=2, keepdim=True) + 1e-8)

    # feat2是feat1的扰动版本
    feat2 = feat1 + 0.15 * torch.randn(B, N, C)
    feat2_norm = feat2 / (feat2.norm(dim=2, keepdim=True) + 1e-8)

    C = 1.0 - torch.bmm(feat1_norm, feat2_norm.transpose(1, 2))

    epsilon = torch.tensor([0.05])
    gamma = torch.tensor([1.0])

    # 逐步可视化Sinkhorn迭代
    K = torch.exp(-C / epsilon)
    max_iter = 5
    power = gamma / (gamma + epsilon + 1e-8)

    a = torch.ones((B, N, 1)) / N
    prob1 = torch.ones((B, N, 1)) / N
    prob2 = torch.ones((B, N, 1)) / N

    fig, axes = plt.subplots(2, 4, figsize=(18, 9))

    # K (初始)
    K_np = K[0].numpy()
    im0 = axes[0, 0].imshow(K_np, cmap='YlOrRd')
    axes[0, 0].set_title('K (exp(-C/ε))\n迭代0')
    axes[0, 0].set_xlabel('Positive'); axes[0, 0].set_ylabel('Query')
    plt.colorbar(im0, ax=axes[0, 0])

    for iteration in range(1, min(max_iter + 1, 7)):
        # Update b
        KTa = torch.bmm(K.transpose(1, 2), a)
        b = torch.pow(prob2 / (KTa + 1e-8), power)
        # Update a
        Kb = torch.bmm(K, b)
        a = torch.pow(prob1 / (Kb + 1e-8), power)

        T = torch.mul(torch.mul(a, K), b.transpose(1, 2))
        T_np = T[0].numpy()

        ax = axes[(iteration) // 4, (iteration) % 4]
        im = ax.imshow(T_np, cmap='YlOrRd')
        ax.set_title(f'Transport Plan T\n迭代{iteration}')
        ax.set_xlabel('Positive'); ax.set_ylabel('Query')
        plt.colorbar(im, ax=ax)

    # 空余位置
    for i in range(max_iter + 1, 8):
        axes[i // 4, i % 4].axis('off')

    plt.suptitle('Sinkhorn 非平衡最优传输迭代过程', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'uot_sinkhorn_iterations.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'  [保存] {path}')


def test_rigid_transform():
    """测试刚体变换估计"""
    print('\n--- Weighted SVD 刚体变换估计 ---')

    torch.manual_seed(42)
    B, N = 2, 150

    # 真实变换
    angle = torch.tensor(0.5)  # ~28.6度
    R_true = torch.tensor([
        [torch.cos(angle), -torch.sin(angle), 0],
        [torch.sin(angle), torch.cos(angle), 0],
        [0, 0, 1]
    ]).unsqueeze(0).repeat(B, 1, 1)
    t_true = torch.tensor([2.0, -1.0, 0.1]).unsqueeze(0).unsqueeze(-1).repeat(B, 1, 1)

    # query点云
    pts1 = torch.randn(B, N, 3) * 20

    # positive点云 = R * query + t + noise
    pts2 = R_true @ pts1.transpose(1, 2) + t_true
    pts2 = pts2.transpose(1, 2) + 0.3 * torch.randn(B, N, 3)

    # 模拟transport weights（前80个点匹配好，后70个匹配差）
    weights = torch.ones(B, N)
    weights[:, 80:] = 0.1  # 降低后70个点的权重

    # 估计变换
    transform = compute_rigid_transform(pts1, pts2, weights)

    # 评估
    R_est = transform[:, :3, :3]
    t_est = transform[:, :3, 3]

    # 旋转误差
    R_err = R_est @ R_true.transpose(1, 2)
    trace = torch.diagonal(R_err, dim1=1, dim2=2).sum(dim=1)
    angle_err = torch.acos(torch.clamp((trace - 1) / 2, -1, 1)) * 180 / np.pi

    # 平移误差
    t_err = (t_est - t_true.squeeze(-1)).norm(dim=1)

    print(f'旋转误差: {angle_err[0].item():.2f}°')
    print(f'平移误差: {t_err[0].item():.3f}m')

    # 可视化
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # 3D点云（XY平面投影）
    pts1_2d = pts1[0, :, :2].numpy()
    pts2_2d = pts2[0, :, :2].numpy()

    # 投影点
    pts1_transformed = (R_est[0] @ pts1[0].T + t_est[0].unsqueeze(-1)).T[:, :2].numpy()

    axes[0].scatter(pts1_2d[:, 0], pts1_2d[:, 1], c='blue', s=10, alpha=0.6, label='Query')
    axes[0].scatter(pts2_2d[:, 0], pts2_2d[:, 1], c='red', s=10, alpha=0.6, label='Positive')
    for i in range(min(20, N)):
        if weights[0, i] > 0.5:
            axes[0].plot([pts1_2d[i, 0], pts2_2d[i, 0]],
                         [pts1_2d[i, 1], pts2_2d[i, 1]],
                         'gray', alpha=0.3, linewidth=0.5)
    axes[0].set_title('匹配点对 (蓝色→红色)')
    axes[0].set_xlabel('X (m)'); axes[0].set_ylabel('Y (m)')
    axes[0].legend(fontsize=8)
    axes[0].set_aspect('equal')

    # 变换后
    axes[1].scatter(pts1_transformed[:, 0], pts1_transformed[:, 1],
                    c='blue', s=10, alpha=0.6, label='Query (变换后)')
    axes[1].scatter(pts2_2d[:, 0], pts2_2d[:, 1],
                    c='red', s=10, alpha=0.6, label='Positive (目标)')
    axes[1].set_title(f'变换后对比\n旋转误差:{angle_err[0].item():.2f}° 平移误差:{t_err[0].item():.3f}m')
    axes[1].set_xlabel('X (m)'); axes[1].set_ylabel('Y (m)')
    axes[1].legend(fontsize=8)
    axes[1].set_aspect('equal')

    plt.suptitle('Weighted SVD 刚体变换估计', fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'uot_rigid_transform.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'  [保存] {path}')


def visualize_epsilon_gamma():
    """可视化epsilon和gamma参数的影响"""
    print('\n--- epsilon/gamma 参数分析 ---')

    torch.manual_seed(42)
    N = 50
    feat1 = torch.randn(1, N, 128)
    feat1_norm = feat1 / (feat1.norm(dim=2, keepdim=True) + 1e-8)
    feat2 = feat1 + 0.2 * torch.randn(1, N, 128)
    feat2_norm = feat2 / (feat2.norm(dim=2, keepdim=True) + 1e-8)

    epsilons = [0.01, 0.05, 0.1, 0.5]
    gammas = [0.1, 1.0, 10.0]

    fig, axes = plt.subplots(len(gammas), len(epsilons), figsize=(16, 12))

    for gi, gamma in enumerate(gammas):
        for ei, eps in enumerate(epsilons):
            epsilon = torch.tensor([eps])
            gam = torch.tensor([gamma])
            T = sinkhorn_unbalanced(
                feat1_norm, feat2_norm,
                epsilon=epsilon, gamma=gam,
                max_iter=5, matrix='cosine'
            )
            ax = axes[gi, ei]
            im = ax.imshow(T[0].numpy(), cmap='YlOrRd')
            ax.set_title(f'ε={eps}, γ={gamma}')
            ax.set_xlabel('Positive'); ax.set_ylabel('Query')
            plt.colorbar(im, ax=ax)

    plt.suptitle('epsilon (熵正则) 和 gamma (质量正则) 对 Transport Plan 的影响',
                 fontsize=14, fontweight='bold')
    plt.tight_layout()
    path = os.path.join(OUTPUT_DIR, 'uot_epsilon_gamma.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f'  [保存] {path}')

    print("""
    参数解释:
    - epsilon (ε): 熵正则化强度
      - 小ε → Transport Plan更稀疏（hard matching）
      - 大ε → Transport Plan更平滑（soft matching）
    - gamma (γ): 质量正则化强度
      - 小γ → 允许部分匹配（质量可增减）
      - 大γ → 要求质量守恒（所有点必须匹配）
    """)


def analyze_parameters():
    """参数量分析"""
    print('\n--- 参数量分析 ---')
    uot = UOTHead(nb_iter=5, name='original')
    total = sum(p.numel() for p in uot.parameters())
    print(f'总参数量: {total} (仅 epsilon, gamma 两个可学习标量)')
    for name, param in uot.named_parameters():
        print(f'  {name}: {param.data.item():.4f}')


def main():
    print('=' * 60)
    print('UOT (Unbalanced Optimal Transport) 位姿估计可视化')
    print('=' * 60)

    analyze_parameters()
    visualize_cost_matrix()
    visualize_sinkhorn()
    test_rigid_transform()
    visualize_epsilon_gamma()

    print('\n' + '=' * 60)
    print('结构总结:')
    print('=' * 60)
    print("""
    UOTHead (非平衡最优传输位姿估计):

    ┌──────────────────────────────────────────────────────┐
    │ 输入: feat1(B,150,128), feat2(B,150,128)              │
    │       kpts1(B,150,3),  kpts2(B,150,3)                │
    │                                                      │
    │ 1. Cost Matrix: C = 1 - cosine_sim(feat1, feat2)     │
    │    → (B, 150, 150)                                   │
    │                                                      │
    │ 2. Sinkhorn Unbalanced OT (迭代5次):                  │
    │    K = exp(-C / epsilon)                             │
    │    for i in range(5):                                │
    │        b = (prob2 / Kᵀa)^(γ/(γ+ε))                   │
    │        a = (prob1 / Kb)^(γ/(γ+ε))                    │
    │    T = a ⊙ K ⊙ bᵀ                                    │
    │    → (B, 150, 150) 运输计划                           │
    │                                                      │
    │ 3. 投影: project_kpts = T @ kpts2 / ΣT               │
    │    → (B, 150, 3)  query匹配点在positive空间的投影坐标   │
    │                                                      │
    │ 4. Weighted SVD 刚体变换:                             │
    │    - 加权中心化                                       │
    │    - SVD分解协方差                                    │
    │    - 输出 R(3×3), t(3×1)                             │
    │    → transformation: (B, 3, 4)                       │
    └──────────────────────────────────────────────────────┘

    为什么用Unbalanced OT（非平衡最优传输）？
    - 标准OT要求两个点集大小相同且质量守恒
    - 实际场景：部分关键点在另一帧中可能被遮挡
    - Unbalanced OT允许部分匹配，更鲁棒

    两个可学习参数:
    - epsilon (ε): 熵正则化，exp(ε)+0.03
    - gamma (γ): 质量正则化，exp(γ)
    """)

    print(f'\n所有可视化结果保存在: {OUTPUT_DIR}')


if __name__ == '__main__':
    main()