import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.datasets import make_circles, make_moons
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

# Same data, same model; only the features change
datasets = {'moons': make_moons(n_samples=400, noise=0.2, random_state=0),
            'circles': make_circles(n_samples=400, noise=0.15, factor=0.5, random_state=0)}
models = {'plain: features x₁, x₂': lambda: make_pipeline(StandardScaler(), LogisticRegression(C=np.inf)),
          'degree 3: x₁, x₂, x₁², x₁x₂, x₂², … (9 features)':
              lambda: make_pipeline(PolynomialFeatures(degree=3, include_bias=False), StandardScaler(),
                                    LogisticRegression(C=np.inf, max_iter=5000))}
cmap = LinearSegmentedColormap.from_list('gb', ['#8b949e', '#58a6ff'])

fig, axes = plt.subplots(2, 2, figsize=(9, 7))
for r, (dname, (X, y)) in enumerate(datasets.items()):
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, stratify=y, random_state=0)
    lo, hi = X.min(axis=0) - 0.4, X.max(axis=0) + 0.4
    g1, g2 = np.meshgrid(np.linspace(lo[0], hi[0], 200), np.linspace(lo[1], hi[1], 200))
    for c, (mname, make) in enumerate(models.items()):
        m = make().fit(X_tr, y_tr)
        ax = axes[r, c]
        P = m.predict_proba(np.c_[g1.ravel(), g2.ravel()])[:, 1].reshape(g1.shape)
        ax.contourf(g1, g2, P, levels=np.linspace(0, 1, 11), cmap=cmap, alpha=0.22, rasterized=True)
        ax.contour(g1, g2, P, levels=[0.5], colors='#f0883e', linewidths=2)
        ax.scatter(*X_te[y_te == 0].T, s=12, color='gray', label='class 0 (test)')
        ax.scatter(*X_te[y_te == 1].T, s=12, color='#58a6ff', label='class 1 (test)')
        ax.set_title(f'{dname}, {mname}\ntrain {m.score(X_tr, y_tr):.1%} · test {m.score(X_te, y_te):.1%}',
                     fontsize=9, loc='left')
        ax.tick_params(labelsize=7.5)
        ax.set_xticks([])
        ax.set_yticks([])
axes[0, 0].legend(fontsize=7.5, loc='lower left')
fig.suptitle('A linear model on curved data: add the curved features, keep the model', fontsize=10)
plt.tight_layout()

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

plt.close()