import numpy as np
import matplotlib.pyplot as plt

# Example sentence
words = ["the", "rain", "in", "Spain", "falls", "mainly", "on", "the", "plain"]
positions = np.arange(len(words))

fig, ax = plt.subplots(figsize=(10, 3))
ax.set_ylim(-1.2, 1.2)
ax.set_xlim(-0.5, len(words) - 0.5)
ax.axis('off')

# Plot words
for x, w in zip(positions, words):
    ax.text(x, 0, w, ha='center', va='center')

# Parameters for curve offset and amplitude
offset_top = 0.4
offset_bottom = -0.4
amplitude = 0.6

# Draw curves for bigrams (2-grams) on top
for i in range(len(words) - 1):
    j = i + 1
    t = np.linspace(0, 1, 100)
    xs = positions[i] + t * (positions[j] - positions[i])
    ys = offset_top + amplitude * np.sin(np.pi * t)
    ax.plot(xs, ys, color='black')

# Draw curves for 1-skip-2-grams on bottom
for i in range(len(words) - 2):
    j = i + 2
    t = np.linspace(0, 1, 100)
    xs = positions[i] + t * (positions[j] - positions[i])
    ys = offset_bottom - amplitude * np.sin(np.pi * t)
    ax.plot(xs, ys, color='black')

plt.tight_layout()
# plt.show()
plt.savefig('skip-gram.svg')