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

# House value by ocean_proximity, groups ordered by their median.
# Left: boxplots (box = Q1..Q3, line = median, whiskers = 1.5 IQR).
# Right: violins of the same groups (cut at the observed min and max).
URL = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
df = pd.read_csv(URL)
df['value'] = df['median_house_value'] / 1000
order = df.groupby('ocean_proximity')['value'].median().sort_values().index
groups = [df.loc[df['ocean_proximity'] == g, 'value'] for g in order]
labels = [f'{g}\nn = {len(v):,}' for g, v in zip(order, groups)]

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, 2, figsize=(10, 3.8), sharey=True)

ax = axes[0]
ax.boxplot(groups, tick_labels=labels, patch_artist=True, widths=0.55,
           boxprops=dict(facecolor=BLUE, alpha=0.6, edgecolor='gray'),
           medianprops=dict(color=ORANGE, lw=2), whiskerprops=dict(color='gray'),
           capprops=dict(color='gray'), flierprops=dict(marker='o', ms=2, mec='gray', alpha=0.4))
ax.set_title('Boxplot: medians and spread side by side', fontsize=9.5)
ax.set_ylabel('median_house_value (US$ thousand)', fontsize=8.5)

ax = axes[1]
parts = ax.violinplot(groups, showmedians=True, showextrema=False, widths=0.8)
for body in parts['bodies']:
    body.set_facecolor(BLUE)
    body.set_edgecolor('gray')
    body.set_alpha(0.6)
parts['cmedians'].set_color(ORANGE)
parts['cmedians'].set_linewidth(2)
ax.set_xticks(range(1, len(order) + 1), labels)
ax.set_title('Violin: the shape, including a bulge at the ceiling', fontsize=9.5)

for ax in axes:
    ax.axhline(500.001, color=RED, ls='--', lw=0.9)
    ax.tick_params(labelsize=7.5)
    ax.set_ylim(0, 550)
axes[1].text(0.6, 512, 'ceiling 500,001', fontsize=7.5, color=RED)
fig.suptitle('House value by distance to the ocean: INLAND stands apart; ISLAND is five districts', fontsize=10)

plt.tight_layout()

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

plt.close()