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

# When is a thin slice enough? The two updates of the lesson's lab step 6, drawn with the same
# seed and in the same order: a 512 x 512 update of rank 4 plus a little noise, then a matrix of
# pure noise (full rank). For each rank r, the error of the best rank-r approximation (truncated
# SVD, Eckart-Young), i.e. the best any B (d x r) times A (r x d) can do, relative to the update.
BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
rng = np.random.default_rng(0)
d = 512
delta = rng.normal(size=(d, 4)) @ rng.normal(size=(4, d)) + 0.1 * rng.normal(size=(d, d))
noise = rng.normal(size=(d, d))


def errors(M, ranks):
    s = np.linalg.svd(M, compute_uv=False)
    total = np.sum(s ** 2)
    return [np.sqrt(np.sum(s[r:] ** 2) / total) for r in ranks]


ranks = np.arange(1, d + 1)
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, ax = plt.subplots(figsize=(8.6, 4.6))
for M, col, name in [(delta, GREEN, 'update of rank 4 + a little noise'), (noise, ORANGE, 'pure noise (full rank)')]:
    e = errors(M, ranks)
    ax.plot(ranks, e, color=col, lw=2.2, label=name)
    marks = [1, 2, 4, 8, 32] if M is delta else [1, 4, 8, 32, 128]
    for r in marks:
        ax.plot(r, e[r - 1], 'o', color=col, ms=6)
        ax.annotate(f'{e[r - 1]:.3f}', (r, e[r - 1]), textcoords='offset points',
                    xytext=(6, 7 if M is noise else -15), fontsize=10.5, color=col)
ax.axvline(4, color='gray', lw=0.8, ls=':')
ax.text(4.3, 0.33, 'r = 4, the true rank', fontsize=10.5, rotation=90, va='bottom')
ax.set_xscale('log', base=2)
ax.set_xticks([1, 2, 4, 8, 16, 32, 64, 128, 256, 512])
ax.set_xticklabels(['1', '2', '4', '8', '16', '32', '64', '128', '256', '512'])
ax.set_xlim(0.8, 600)
ax.set_ylim(-0.03, 1.08)
ax.set_xlabel('rank r of the slice B·A   (it trains 2 × 512 × r numbers; the full update has 262,144)')
ax.set_ylabel('error of the best rank-r slice')
for side in ('top', 'right'):
    ax.spines[side].set_visible(False)
for side in ('left', 'bottom'):
    ax.spines[side].set_color('gray')
ax.legend(frameon=False, loc='center right', labelcolor='gray', fontsize=11)
plt.tight_layout()

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

plt.close()