import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor

# Rent (R$ thousand) against area (m²): 60 training apartments between 30 and
# 120 m², and fresh apartments up to 200 m² that no model saw. The truth is a
# straight line plus noise. A regression tree predicts the mean of each leaf:
# a step function, flat beyond the last training point.
rng = np.random.default_rng(0)
truth = lambda a: 0.3 + 0.028 * a
area = rng.uniform(30, 120, 60)
rent = truth(area) + rng.normal(0, 0.25, 60)
new_in = rng.uniform(30, 120, 200)
new_out = rng.uniform(120, 200, 200)
y_in = truth(new_in) + rng.normal(0, 0.25, 200)
y_out = truth(new_out) + rng.normal(0, 0.25, 200)

models = {'tree, depth 1': DecisionTreeRegressor(max_depth=1),
          'tree, depth 3': DecisionTreeRegressor(max_depth=3),
          'linear regression': LinearRegression()}
colors = {'tree, depth 1': '#3fb950', 'tree, depth 3': '#f0883e', 'linear regression': '#58a6ff'}
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, ax = plt.subplots(figsize=(8.6, 4.0))
ax.axvspan(120, 205, color='gray', alpha=0.1, lw=0)
ax.text(162, 0.85, 'outside the training range\n(no apartment above 120 m²)', ha='center', fontsize=8)
ax.scatter(new_out, y_out, s=10, facecolors='none', edgecolors='gray', lw=0.6, alpha=0.6,
           label='new apartments, 120–200 m²')
ax.scatter(area, rent, s=16, color='gray', label='training apartments')
grid = np.linspace(25, 205, 1000).reshape(-1, 1)
for name, m in models.items():
    m.fit(area.reshape(-1, 1), rent)
    rmse_in = np.sqrt(np.mean((m.predict(new_in.reshape(-1, 1)) - y_in) ** 2))
    rmse_out = np.sqrt(np.mean((m.predict(new_out.reshape(-1, 1)) - y_out) ** 2))
    ax.plot(grid, m.predict(grid), color=colors[name], lw=2,
            label=f'{name}: RMSE {rmse_in:.2f} inside · {rmse_out:.2f} outside')
top = models['tree, depth 3'].predict([[200]])[0]
ax.annotate(f'every tree prediction above 120 m²\nis the last leaf: R$ {top:.2f} thousand',
            xy=(185, top), xytext=(126, 6.55), fontsize=8,
            arrowprops=dict(arrowstyle='->', color='gray'))
ax.set_xlim(25, 205)
ax.set_ylim(0.5, 7.4)
ax.set_xlabel('area (m²)')
ax.set_ylabel('rent (R$ thousand / month)')
ax.set_title('A regression tree is a step function, and it cannot extrapolate', fontsize=10)
ax.legend(fontsize=7.5, loc='upper left', frameon=False)
ax.spines[['top', 'right']].set_visible(False)
plt.tight_layout()

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

plt.close()