import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

# gray text and axes so the figure reads on light and dark backgrounds
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The worked example of the lesson: 10 subscribers, one new one, and the circles
# that contain its 1, 3 and 5 nearest neighbors
names = ['Ana', 'Bruno', 'Carla', 'Davi', 'Elisa', 'Felipe', 'Gabi', 'Hugo', 'Iara', 'João']
action = np.array([5, 2, 6, 4, 7, 1, 8, 1, 9, 1])
docs = np.array([4, 5, 6, 7, 3, 2, 6, 8, 2, 9])
liked = np.array([1, 0, 0, 1, 1, 0, 1, 0, 1, 0])
q = np.array([4, 4])

d = np.sqrt((action - q[0]) ** 2 + (docs - q[1]) ** 2)
order = np.argsort(d)
BLUE, ORANGE, GREEN, GRAY = '#58a6ff', '#f0883e', '#3fb950', 'gray'

fig, ax = plt.subplots(figsize=(7.2, 5.6))
lines = []
for k, style, angle, off in [(1, '-', 250, 0.35), (3, '--', 115, -0.4), (5, ':', 285, 0.35)]:
    r = d[order[k - 1]]
    votes = liked[order[:k]].sum()
    verdict = 'liked' if votes * 2 > k else 'not liked'
    ax.add_patch(plt.Circle(q, r, fill=False, ls=style, lw=1.4, color=GRAY))
    a = np.deg2rad(angle)
    ax.text(q[0] + (r + off) * np.cos(a), q[1] + (r + off) * np.sin(a), f'k={k}', fontsize=8.5, color=GRAY,
            ha='center', va='center')
    lines.append(f'k = {k}: {votes} of {k} neighbors liked it  →  {verdict}')
ax.text(10.2, -1.2, '\n'.join(lines), fontsize=8.5, color=GRAY, ha='right', va='bottom')
for cls, color, label in [(1, BLUE, 'liked the series'), (0, ORANGE, 'did not like it')]:
    m = liked == cls
    ax.scatter(action[m], docs[m], s=70, color=color, label=label, zorder=3)
for i, n in enumerate(names):
    ax.annotate(f'{n} ({d[i]:.2f})', (action[i], docs[i]), xytext=(6, 5), textcoords='offset points',
                fontsize=8, color=GRAY)
ax.scatter(*q, marker='*', s=260, color=GREEN, zorder=4, label='new subscriber (4, 4)')
ax.set_xlim(-2.5, 10.5)
ax.set_ylim(-1.5, 10)
ax.set_aspect('equal')
ax.set_xlabel('hours of action films per week')
ax.set_ylabel('hours of documentaries per week')
ax.legend(fontsize=8, loc='upper right', frameon=False)
ax.set_title('The same query, three values of k, three circles (distance in parentheses)', fontsize=10)
ax.tick_params(labelsize=8)

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()