import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap

# 20,640 districts of California housing: income against house value, drawn
# as a plain scatter, as a scatter with transparency, and as a hexbin.
# The scatter layers are rasterized so the SVG stays small.
URL = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
df = pd.read_csv(URL)
x, y = df['median_income'], df['median_house_value'] / 1000

BLUE, ORANGE, RED = '#58a6ff', '#f0883e', '#ff7b72'
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.5), sharex=True, sharey=True)

axes[0].scatter(x, y, s=4, color=BLUE, lw=0, rasterized=True)
axes[0].set_title('plain scatter: a solid blob', fontsize=9.5)

axes[1].scatter(x, y, s=4, color=BLUE, alpha=0.08, lw=0, rasterized=True)
axes[1].set_title('alpha = 0.08: density appears', fontsize=9.5)
capped = (df['median_house_value'] >= 500001).sum()
axes[1].annotate(f'the ceiling: {capped} districts\nat exactly 500,001', xy=(11.5, 500), xytext=(7.0, 560),
                 fontsize=7.5, color=RED, arrowprops=dict(arrowstyle='->', color=RED, lw=0.8))

cmap = LinearSegmentedColormap.from_list('blue_orange', [BLUE, ORANGE, RED])
hb = axes[2].hexbin(x, y, gridsize=40, mincnt=1, cmap=cmap, linewidths=0)
axes[2].set_title('hexbin: districts per cell', fontsize=9.5)
cb = fig.colorbar(hb, ax=axes[2], fraction=0.05, pad=0.02)
cb.set_label('districts per cell', fontsize=8)
cb.ax.tick_params(labelsize=7.5)
cb.outline.set_edgecolor('gray')

for ax in axes:
    ax.set_xlabel('median_income (US$ 10k)', fontsize=8.5)
    ax.tick_params(labelsize=8)
    ax.set_ylim(0, 620)
axes[0].set_ylabel('median_house_value (US$ thousand)', fontsize=8.5)
fig.suptitle(f'{len(df):,} points on one chart: fix overplotting before reading anything', fontsize=10)

plt.tight_layout()

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

plt.close()