import time
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import make_classification
from sklearn.linear_model import SGDClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC, LinearSVC

# Fit time against the number of training rows (20 features, 5% label noise).
# Absolute times depend on the machine; the slopes are what matter.
X_all, y_all = make_classification(n_samples=256_000, n_features=20, n_informative=8,
                                   flip_y=0.05, random_state=0)
X_all = StandardScaler().fit_transform(X_all)


def fit_time(model, n, repeats=3):
    best = np.inf
    for _ in range(repeats):
        t = time.perf_counter()
        model.fit(X_all[:n], y_all[:n])
        best = min(best, time.perf_counter() - t)
    return best


ns_kernel = [1000 * 2 ** k for k in range(6)]            # up to 32,000
ns_linear = [1000 * 2 ** k for k in range(9)]            # up to 256,000
models = [
    ('SVC, RBF kernel', SVC(kernel='rbf'), ns_kernel, '#f0883e'),
    ('LinearSVC', LinearSVC(), ns_linear, '#58a6ff'),
    ("SGDClassifier(loss='hinge')", SGDClassifier(loss='hinge', random_state=0), ns_linear, '#3fb950'),
]

fig, ax = plt.subplots(figsize=(7.2, 4.0))
slopes = {}
for name, model, ns, color in models:
    t = np.array([fit_time(model, n, repeats=1 if (n > 8000 and isinstance(model, SVC)) else 3)
                  for n in ns])
    slope = np.polyfit(np.log(ns[-4:]), np.log(t[-4:]), 1)[0]   # slope on the largest sizes
    slopes[name] = slope
    ax.loglog(ns, t, 'o-', color=color, ms=4, lw=1.6,
              label=f'{name}: time ∝ n$^{{{slope:.1f}}}$, {t[-1]:.2f} s at n = {ns[-1]:,}')
    if isinstance(model, SVC):
        # where the kernel SVM would be at 256,000 rows if the slope held
        t_far = t[-1] * (256_000 / ns[-1]) ** slope
        ax.loglog([ns[-1], 256_000], [t[-1], t_far], ':', color=color, lw=1.4)
        ax.annotate(f'same slope continued:\n≈ {t_far:,.0f} s at 256,000', xy=(256_000, t_far),
                    xytext=(70_000, t_far / 25), fontsize=7.5, color=color, va='center')
        ax.annotate(f'{model.n_support_.sum():,} support vectors\n(of {ns[-1]:,} rows)',
                    xy=(ns[-1], t[-1]), xytext=(ns[-1] * 1.5, t[-1] / 8), fontsize=7.5, color='gray',
                    arrowprops=dict(arrowstyle='->', color='gray'))
ax.set_xlabel('training rows n', fontsize=9)
ax.set_ylabel('fit time (s), log scale', fontsize=9)
ax.tick_params(labelsize=8)
ax.grid(True, which='major', color='gray', alpha=0.2)
ax.legend(fontsize=7.5, loc='upper left', framealpha=0.3)
k, l1, l2 = slopes.values()
ax.set_title(f'Doubling n multiplies the fit time by {2 ** k:.1f} for the kernel SVM, '
             f'by {2 ** l2:.1f}–{2 ** l1:.1f} for the linear ones', fontsize=9.5)

plt.tight_layout()

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

plt.close()