# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from mpl_toolkits.mplot3d import Axes3D  # 注册 3D 投影

# —— 1. 读取坐标和元数据 —— #
df = pd.read_csv('umap_3d_points.csv')

# —— 2. 划分三类 —— #
idx1     = df.index[df['label'] == 1].tolist()
idx0     = df.index[df['label'] == 0].tolist()
idx_rest = df.index[df['label'].isna()].tolist()

# —— 3. 渐变色 —— #
cmap = LinearSegmentedColormap.from_list('bg_cmap', ['#87CEEB', '#FFDAB9'])
bg_scores = df.loc[idx_rest, 'score'].fillna(0.5).values
bg_colors = cmap(bg_scores)

# —— 4. 绘图 —— #
fig = plt.figure(figsize=(8,8))
ax = fig.add_subplot(111, projection='3d')

# 背景点
ax.scatter(
    df.loc[idx_rest,'x'],
    df.loc[idx_rest,'y'],
    df.loc[idx_rest,'z'],
    c=bg_colors,
    s=20,
    alpha=0.06,
    label='Background'
)

# 标签 0
ax.scatter(
    df.loc[idx0,'x'],
    df.loc[idx0,'y'],
    df.loc[idx0,'z'],
    c='#006EFF',
    s=60,
    alpha=1.0,
    label='Not Improved'
)

# 标签 1
ax.scatter(
    df.loc[idx1,'x'],
    df.loc[idx1,'y'],
    df.loc[idx1,'z'],
    c='orange',
    s=60,
    alpha=1.0,
    label='Improved'
)

ax.set_title('3D UMAP Visualization', fontsize=18)
# ax.set_axis_off()
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
ax.legend(loc='upper left', frameon=False, fontsize=12)
plt.tight_layout()
plt.show()  # 或者显示
# plt.savefig('umap_3d.svg')
# print("已保存 umap_3d.svg")