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

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The five points of the worked example
X = np.array([[2, 4], [4, 2], [5, 5], [6, 8], [8, 6]], dtype=float)
mean = X.mean(axis=0)
Xc = X - mean
C = Xc.T @ Xc / (len(X) - 1)                  # [[5, 3], [3, 5]]
lam, V = np.linalg.eigh(C)
lam, V = lam[::-1], V[:, ::-1]                # largest first: 8, 2
w1, w2 = V[:, 0], V[:, 1]
if w1[0] < 0:
    w1 = -w1
if w2[1] < 0:
    w2 = -w2
X_rec = mean + np.outer(Xc @ w1, w1)          # keep PC1 only

fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 4.0), gridspec_kw={'width_ratios': [1, 1.25]})

# left: the cloud, the two components, the projections onto PC1 and the residuals
t = np.linspace(-5, 5, 2)
left.plot(mean[0] + t * w1[0], mean[1] + t * w1[1], color='#58a6ff', lw=1, alpha=0.6)
left.plot(mean[0] + t * w2[0], mean[1] + t * w2[1], color='#f0883e', lw=1, alpha=0.6)
for p, r in zip(X, X_rec):
    left.plot([p[0], r[0]], [p[1], r[1]], color='#ff7b72', lw=1.5, ls=':')
left.scatter(X[:, 0], X[:, 1], s=45, color='gray', zorder=3, label='data')
left.scatter(X_rec[:, 0], X_rec[:, 1], s=45, facecolors='none', edgecolors='#58a6ff', lw=1.5,
             zorder=3, label='reconstruction with PC1 only')
left.scatter(*mean, marker='+', s=120, color='#3fb950', zorder=4, label='mean (5, 5)')
left.annotate('', mean + np.sqrt(lam[0]) * w1, mean,
              arrowprops=dict(arrowstyle='->', color='#58a6ff', lw=2.2))
left.annotate('', mean + np.sqrt(lam[1]) * w2, mean,
              arrowprops=dict(arrowstyle='->', color='#f0883e', lw=2.2))
left.text(*(mean + np.sqrt(lam[0]) * w1 + [0.15, -0.55]), r'PC1, $\lambda_1 = 8$', color='#58a6ff', fontsize=9)
left.text(*(mean + np.sqrt(lam[1]) * w2 + [-2.2, 0.25]), r'PC2, $\lambda_2 = 2$', color='#f0883e', fontsize=9)
left.set_xlim(0.5, 9.5)
left.set_ylim(0.5, 9.5)
left.set_aspect('equal')
left.set_xlabel('x', fontsize=9)
left.set_ylabel('y', fontsize=9)
left.tick_params(labelsize=8)
left.legend(fontsize=7.5, loc='upper left', framealpha=0)
left.set_title('dotted red = what PC1 throws away', fontsize=9.5, loc='left')

# right: variance captured and residual as the line turns
theta = np.linspace(0, 180, 361)
w = np.column_stack([np.cos(np.radians(theta)), np.sin(np.radians(theta))])
captured = np.einsum('ij,jk,ik->i', w, C, w)          # w^T C w
residual = np.trace(C) - captured
right.plot(theta, captured, color='#58a6ff', lw=2, label=r'variance captured $w^\top C w$')
right.plot(theta, residual, color='#ff7b72', lw=2, label='variance thrown away')
right.axhline(np.trace(C), color='gray', ls='--', lw=1)
right.text(2, np.trace(C) + 0.2, 'sum = trace = 10 at every angle', color='gray', fontsize=8.5)
right.axvline(45, color='#58a6ff', ls=':', lw=1)
right.annotate('45°: 8 kept, 2 lost', (45, 8), textcoords='offset points', xytext=(8, -4),
               color='#58a6ff', fontsize=8.5)
right.set_xlabel('angle of the line (degrees)', fontsize=9)
right.set_ylabel('variance', fontsize=9)
right.set_xticks([0, 45, 90, 135, 180])
right.set_ylim(0, 11.5)
right.tick_params(labelsize=8)
right.legend(fontsize=8, loc='lower right', framealpha=0)
right.set_title('one optimization, two readings', fontsize=9.5, loc='left')

plt.tight_layout()

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

plt.close()