import matplotlib.pyplot as plt
import pandas as pd
from mpl_ornaments.titles import set_title_and_subtitle

#1
grid_spec = {'hspace':0.35, 'wspace':0.3}
fig, axs = plt.subplots(nrows=4, ncols=4, figsize=(8.5,9), 
                        gridspec_kw=grid_spec)
axs = axs.flatten()

df = pd.read_csv('data/datasaurus_dozen.csv', comment='#')

#2
cols, funs = ('x', 'y'), ('mean', 'std')
stats = {f'{col} {fun}': df[col].agg(fun) for col in cols for fun in funs}
stats['r (x,y)'] = df['x'].corr(df['y'])

#3
text = str().join([f'{key} = {val:4.2f}\n' for key, val in stats.items()])
axs[0].text(0.15, 0.45, s=text, va='center')
axs[0].set_title('Summary stats')
axs[0].set_xticks([]); axs[0].set_yticks([])

#4
groups = df.groupby('dataset')
for idx, (name, group) in enumerate(groups):
    axs[idx+1].set_title(name)
    axs[idx+1].scatter(x=group['x'], y=group['y'], edgecolors='black', 
                       linewidths=0.5)

#5
[ax.axis('off') for ax in axs[-2::]] 

title = 'The Datasaurus dozen'
subtitle = 'Data generated with R ver. 4.1.2 and datasauRus ver. 0.1.7.'
set_title_and_subtitle(
    fig=fig, title=title, subtitle=subtitle, alignment='left', 
    v_offset=10, h_offset=55)

fig.savefig('charts/scatter-plot-datasaurus-dozen.png', 
            bbox_inches='tight', dpi=300)