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

# Two right-skewed columns of California housing before and after a log.
# Skewness is pandas' .skew(); the red bars are the districts with fewer
# than 50 inhabitants, which the log turns into a new left tail.
URL = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
df = pd.read_csv(URL)

BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(2, 2, figsize=(9.5, 5.2))

for row, (col, label) in enumerate([('median_income', 'median_income (US$ 10k)'),
                                    ('population', 'population (inhabitants)')]):
    x = df[col]
    tiny = df['population'] < 50
    for c, (values, name) in enumerate([(x, label), (np.log(x), f'log({col})')]):
        ax = axes[row, c]
        raw_pop = col == 'population' and c == 0
        edges = np.histogram_bin_edges(values, bins=60, range=(0, 8000) if raw_pop else None)
        fmt = ',.0f' if raw_pop else ',.2f'
        ax.hist(values, bins=edges, color=BLUE, edgecolor='none')
        if col == 'population':
            ax.hist(values[tiny], bins=edges, color=RED, edgecolor='none',
                    label=f'{tiny.sum()} districts with < 50 people')
        ax.axvline(values.mean(), color=ORANGE, lw=1.5, label=f'mean {values.mean():{fmt}}')
        ax.axvline(values.median(), color=GREEN, lw=1.5, ls='--', label=f'median {values.median():{fmt}}')
        ax.set_title(f'{name}: skewness {values.skew():+.2f}', fontsize=9.5)
        ax.set_xlabel(name, fontsize=8.5)
        ax.tick_params(labelsize=8)
        ax.set_yticks([])
        ax.legend(fontsize=7.5, frameon=False, loc='upper right' if c == 0 else 'upper left')
axes[1, 0].set_xlim(0, 8000)
axes[1, 0].text(0.97, 0.35, 'x-axis cut at 8,000\n(max is 35,682)', transform=axes[1, 0].transAxes,
                ha='right', fontsize=7.5)
axes[1, 1].annotate('the new left tail:\nsmall districts that\nwere invisible above', xy=(3.2, 12), xytext=(1.3, 450),
                    fontsize=7.5, color=RED, arrowprops=dict(arrowstyle='->', color=RED, lw=0.8))
fig.suptitle('The log pulls in a long right tail, and can overshoot into a left one', fontsize=10)

plt.tight_layout()

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

plt.close()