import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import load_wine
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import pairwise_distances
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

# Left: relative contrast (farthest - nearest) / nearest for 1000 uniform points, as the
# dimension grows (median over 100 query points).
# Right: wine + pure-noise columns; mean hold-out accuracy over 20 draws of noise and split.
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
BLUE, ORANGE = '#58a6ff', '#f0883e'
dims = [2, 3, 5, 10, 20, 50, 100, 200, 500, 1000]
contrast = []
for d in dims:
    rng = np.random.default_rng(d)                  # same draws as the lesson's code block
    P = rng.uniform(size=(1000, d))
    Q = rng.uniform(size=(100, d))
    D = pairwise_distances(Q, P)
    contrast.append(np.median((D.max(axis=1) - D.min(axis=1)) / D.min(axis=1)))

Xw, yw = load_wine(return_X_y=True)
noise = [0, 3, 10, 30, 100, 300, 1000]
acc = {'k-NN (k = 5)': [], 'logistic regression': []}
for n in noise:
    runs = {name: [] for name in acc}
    for seed in range(20):
        r = np.random.default_rng(seed)
        Xn = np.hstack([Xw, r.normal(size=(len(Xw), n))])
        a, b, c, d = train_test_split(Xn, yw, test_size=0.3, stratify=yw, random_state=seed)
        runs['k-NN (k = 5)'].append(
            make_pipeline(StandardScaler(), KNeighborsClassifier(5)).fit(a, c).score(b, d))
        runs['logistic regression'].append(
            make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000)).fit(a, c).score(b, d))
    for name in acc:
        acc[name].append(np.mean(runs[name]))

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 3.9))
left.plot(dims, contrast, 'o-', color=BLUE, ms=4)
for d, c in zip(dims, contrast):
    if d in (2, 10, 100, 1000):
        left.annotate(f'{c:.2f}' if c < 10 else f'{c:.0f}', (d, c), xytext=(5, 5),
                      textcoords='offset points', fontsize=8)
left.set_xscale('log')
left.set_yscale('log')
left.set_xlabel('number of dimensions d')
left.set_ylabel('(farthest − nearest) / nearest\nmedian of 100 queries')
left.set_title('Distances concentrate: nearest ≈ farthest', fontsize=10)
left.tick_params(labelsize=8)

x = np.array(noise) + 1
for (name, vals), color in zip(acc.items(), [BLUE, ORANGE]):
    right.plot(x, vals, 'o-', color=color, ms=4, label=name)
    right.annotate(f'{vals[-1]:.0%}', (x[-1], vals[-1]), xytext=(4, -3), textcoords='offset points',
                   fontsize=8, color=color)
right.axhline(71 / 178, color='gray', lw=0.8, ls=':')
right.text(2000, 71 / 178 + 0.015, 'majority class (40%)', fontsize=8, ha='right')
right.set_xscale('log')
right.set_xticks(x, [str(n) for n in noise])
right.set_xlabel('pure-noise columns appended to the 13 wine features')
right.set_ylabel('test accuracy (mean of 20 draws)')
right.set_ylim(0.3, 1.03)
right.set_xlim(0.8, 2300)
right.legend(fontsize=8, loc='center left', frameon=False)
right.set_title('Irrelevant features drown the neighbors', fontsize=10)
right.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()