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

# Hardcoded 6-dimensional "embeddings" for 5 words
words = ["cat", "dog", "car", "bus", "apple"]
embeddings = np.array([
    [0.4, 0.2, 0.5, 0.1, 0.4, 0.2],  # cat
    [0.5, 0.1, 0.4, 0.2, 0.3, 0.2],  # dog
    [0.1, 0.6, 0.2, 0.8, 0.2, 0.3],  # car
    [0.2, 0.7, 0.1, 0.9, 0.1, 0.3],  # bus
    [0.8, 0.2, 0.9, 0.1, 0.7, 0.2],  # apple
])

# PCA to 2D
pca = PCA(n_components=2)
reduced_pca = pca.fit_transform(embeddings)

# t-SNE to 2D (small perplexity for small dataset)
tsne = TSNE(n_components=2, random_state=42, perplexity=2)
reduced_tsne = tsne.fit_transform(embeddings)

# Plotting
def plot_embeddings(X, title):
    plt.figure(figsize=(5, 4))
    for i, word in enumerate(words):
        x, y = X[i]
        plt.scatter(x, y)
        plt.text(x+0.01, y+0.01, word, fontsize=12)
    plt.title(title)
    plt.axis('off')
    plt.show()

plot_embeddings(reduced_pca, "PCA Visualization")
plot_embeddings(reduced_tsne, "t-SNE Visualization")