import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.metrics import pairwise_distances

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


def relative_contrast(X):
    """For each point: (farthest - nearest) / nearest. Median over the points."""
    D = pairwise_distances(X)
    np.fill_diagonal(D, np.nan)                 # a point is not its own neighbor
    nearest, farthest = np.nanmin(D, axis=1), np.nanmax(D, axis=1)
    return np.median((farthest - nearest) / nearest)


# Same recipe as the lesson's code block (one seed per d, so the numbers match)
dims = [2, 3, 5, 10, 20, 50, 100, 200, 500, 1000]
iid, plane = [], []
for d in dims:
    rng = np.random.default_rng(d)
    noise = rng.random((400, d))                                   # d independent columns
    sheet = np.hstack([rng.random((400, 2)), np.zeros((400, d - 2))])
    Q, _ = np.linalg.qr(rng.normal(size=(d, d)))                   # a random rotation
    iid.append(relative_contrast(noise))
    plane.append(relative_contrast(sheet @ Q))                     # 2 real degrees of freedom

fig, ax = plt.subplots(figsize=(7.5, 3.8))
ax.plot(dims, iid, 'o-', color='#f0883e', lw=2,
        label='d independent uniform columns')
ax.plot(dims, plane, 's-', color='#58a6ff', lw=2,
        label='a 2-D sheet rotated into d dimensions')
tail = np.array([20, 1000])
ax.plot(tail, iid[-1] * np.sqrt(1000 / tail) * 1.6, '--', color='gray', lw=1,
        label=r'slope of $1/\sqrt{d}$, for reference')
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('number of columns d')
ax.set_ylabel(r'relative contrast $(d_{max}-d_{min})/d_{min}$')
ax.annotate(f'{iid[0]:.1f}', (2, iid[0]), textcoords='offset points', xytext=(8, -14),
            color='#f0883e', fontsize=9)
ax.annotate(f'{iid[-1]:.2f}', (1000, iid[-1]), textcoords='offset points', xytext=(8, -3),
            color='#f0883e', fontsize=9)
ax.annotate(f'{plane[-1]:.1f}', (1000, plane[-1]), textcoords='offset points', xytext=(8, -3),
            color='#58a6ff', fontsize=9)
ax.set_xlim(1.6, 2200)
ax.set_title('400 points: nearest vs farthest neighbor, median over the points', fontsize=10, loc='left')
ax.legend(fontsize=8, loc='lower left', framealpha=0)
ax.tick_params(labelsize=8)
ax.grid(alpha=0.25, which='both')
plt.tight_layout()

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

plt.close()