import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier

# The true boundary is the diagonal x2 = x1. 300 training points, accuracy
# measured on 5,000 fresh points from the same distribution.
rng = np.random.default_rng(1)
X = rng.uniform(0, 1, (300, 2)); y = (X[:, 1] > X[:, 0]).astype(int)
X_new = rng.uniform(0, 1, (5000, 2)); y_new = (X_new[:, 1] > X_new[:, 0]).astype(int)
add_diff = lambda Z: np.c_[Z, Z[:, 1] - Z[:, 0]]          # engineered feature x2 - x1

models = [('tree on x₁, x₂', DecisionTreeClassifier(random_state=0), lambda Z: Z),
          ('logistic regression on x₁, x₂', LogisticRegression(C=100), lambda Z: Z),
          ('tree on x₁, x₂ and x₂ − x₁', DecisionTreeClassifier(random_state=0), add_diff)]
xx, yy = np.meshgrid(np.linspace(0, 1, 300), np.linspace(0, 1, 300))
grid = np.c_[xx.ravel(), yy.ravel()]

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 3, figsize=(10, 3.6))
for ax, (name, m, f) in zip(axes, models):
    m.fit(f(X), y)
    acc = m.score(f(X_new), y_new)
    Z = m.predict(f(grid)).reshape(xx.shape)
    ax.contourf(xx, yy, Z, levels=[-0.5, 0.5, 1.5], colors=[ORANGE, BLUE], alpha=0.2).set_rasterized(True)
    ax.plot([0, 1], [0, 1], color='gray', lw=1, ls='--')
    ax.scatter(X[y == 0, 0], X[y == 0, 1], s=6, color=ORANGE)
    ax.scatter(X[y == 1, 0], X[y == 1, 1], s=6, color=BLUE)
    size = f'{m.get_n_leaves()} leaves' if hasattr(m, 'get_n_leaves') else '3 coefficients'
    ax.set_title(f'{name}\n{size} · accuracy on new points {acc:.3f}', fontsize=8.5)
    ax.set_xlim(0, 1); ax.set_ylim(0, 1)
    ax.set_xlabel('x₁'); ax.set_aspect('equal')
axes[0].set_ylabel('x₂')
plt.tight_layout()

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

plt.close()