import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn.manifold import TSNE
from matplotlib.lines import Line2D

contras_embeddings = np.load('/ossfs/workspace/contras-codegen-embeddings.npy')
magpie_embeddings = np.load('/ossfs/workspace/magpie-codegen-few-shot-embeddings.npy')
reference_embeddings = np.load('/ossfs/workspace/lcb-codegen-embeddings.npy')

print(contras_embeddings.shape) 
print(magpie_embeddings.shape)
print(reference_embeddings.shape)

embeddings = np.vstack((contras_embeddings, magpie_embeddings, reference_embeddings))


labels = np.array([0] * contras_embeddings.shape[0] + [1] * magpie_embeddings.shape[0] + [2] * reference_embeddings.shape[0])

tsne = TSNE(n_components=2, random_state=42)
embeddings_2d = tsne.fit_transform(embeddings)

plt.figure(figsize=(8, 6))


colors = ['#1f77b4', '#f3d19c', '#E4080A']
labels_text = ["DOMINO", "MAGPIE-Few Shot", "Reference"]
sizes = [2, 2, 10]  

for i in range(3):
    plt.scatter(
        embeddings_2d[labels == i, 0],
        embeddings_2d[labels == i, 1],
        c=colors[i],
        label=labels_text[i],
        s=sizes[i],
        alpha=0.3
    )

legend_handles = [
    Line2D([0], [0], marker='o', color='w', label=labels_text[i],
           markerfacecolor=colors[i], markersize=10)  # 控制 legend 圆点大小
    for i in range(3)
]

plt.legend(handles=legend_handles, loc="upper center", bbox_to_anchor=(0.5, 1.12), ncol=3, fontsize=18, frameon=False, columnspacing=0.5)
# plt.legend(handles, ["DOMINO", "MAGPIE-Few Shot", "Reference"], loc="upper center", bbox_to_anchor=(0.5, 1.2), ncol=3, fontsize=18, frameon=False)

# plt.title("t-SNE Visualization of Embeddings")
plt.tick_params(axis='both', labelsize=12)
plt.savefig("/ossfs/workspace/tsne_visualization-gen1.pdf", format="pdf", dpi=150, bbox_inches="tight")
plt.savefig("/ossfs/workspace/tsne_visualization-gen1.png", format="png", dpi=150, bbox_inches="tight")
plt.show()
