import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA


# plot embedding vectors using PCA
def plot_embedding_pca(embedding_vectors, labels=None):
    # embedding_vectors: list of embedding vectors
    # labels: list of labels (e.g. token name) for each embedding vector

    if labels is None:
        labels = range(len(embedding_vectors))
    reduced_vectors = PCA(n_components=2, random_state=42).fit_transform(
        embedding_vectors
    )  # Apply dimensionality reduction

    # Create the plot
    plt.figure(figsize=(10, 6))
    origin = np.zeros((len(embedding_vectors), 2))
    for i, label in enumerate(labels):
        plt.quiver(
            origin[:, 0],
            origin[:, 1],
            reduced_vectors[i, 0],
            reduced_vectors[i, 1],
            angles="xy",
            scale_units="xy",
            scale=1,
            width=0.002,
        )
        plt.scatter(reduced_vectors[i, 0], reduced_vectors[i, 1])
        plt.annotate(label, (reduced_vectors[i, 0], reduced_vectors[i, 1]))

    plt.title("Embedding Visualization (using 2 Principal Components)")
    plt.show()
