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

# Simpson's paradox in California housing. Left: the share of rooms that are
# bedrooms against house value, all districts, with the least-squares line.
# Right: the same districts split into five income bands (quintiles of
# median_income), one line per band. Lines and r use every district; the dots
# are a random sample of 3,000 so the picture stays readable.
URL = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
df = pd.read_csv(URL).dropna(subset=['total_bedrooms'])
df['bedrooms_per_room'] = df['total_bedrooms'] / df['total_rooms']
df['value'] = df['median_house_value'] / 1000
df['band'] = pd.qcut(df['median_income'], 5, labels=['Q1', 'Q2', 'Q3', 'Q4', 'Q5'])
sample = df.sample(3000, random_state=0)

BLUE, ORANGE, GREEN, RED, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#bc8cff'
colors = dict(zip(['Q1', 'Q2', 'Q3', 'Q4', 'Q5'], [BLUE, GREEN, ORANGE, RED, PURPLE]))
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(10, 3.9), sharex=True, sharey=True)


def fit_line(ax, d, color, label, lw=2.2):
    lo, hi = d['bedrooms_per_room'].quantile([0.02, 0.98])
    b, a = np.polyfit(d['bedrooms_per_room'], d['value'], 1)
    xs = np.linspace(lo, hi, 20)
    ax.plot(xs, a + b * xs, color=color, lw=lw, label=label)


ax = axes[0]
ax.scatter(sample['bedrooms_per_room'], sample['value'], s=5, color='gray', alpha=0.35, lw=0, rasterized=True)
r_all = df['bedrooms_per_room'].corr(df['value'])
fit_line(ax, df, ORANGE, f'all districts: r = {r_all:.2f}', lw=2.6)
ax.set_title('All districts: more bedrooms per room, cheaper', fontsize=9.5)
ax.legend(fontsize=8, frameon=False, loc='upper right')
ax.set_ylabel('median_house_value (US$ thousand)', fontsize=8.5)

ax = axes[1]
for band, d in df.groupby('band', observed=True):
    s = sample[sample['band'] == band]
    ax.scatter(s['bedrooms_per_room'], s['value'], s=5, color=colors[band], alpha=0.3, lw=0, rasterized=True)
for band, d in df.groupby('band', observed=True):
    r = d['bedrooms_per_room'].corr(d['value'])
    fit_line(ax, d, colors[band], f'income {band} (median {d["median_income"].median():.1f}): r = {r:+.2f}')
ax.set_title('Within each income band: the sign flips in 4 of 5', fontsize=9.5)
ax.legend(fontsize=7.5, frameon=False, loc='upper right')

for ax in axes:
    ax.set_xlim(0.1, 0.4)
    ax.set_ylim(0, 540)
    ax.set_xlabel('bedrooms_per_room = total_bedrooms / total_rooms', fontsize=8.5)
    ax.tick_params(labelsize=8)
fig.suptitle("Simpson's paradox: income drives both variables, and hides the within-band trend", fontsize=10)

plt.tight_layout()

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

plt.close()