import matplotlib.pyplot as plt
import pandas as pd
from scipy.stats import pearsonr, spearmanr 
from mpl_ornaments.titles import set_title_and_subtitle

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5,6))

#1
df = pd.read_csv('data/iris.csv', comment='#')
x, y = df['SepalLengthCm'], df['PetalLengthCm']  

#2
ax.scatter(x=x, y=y, edgecolors='black', linewidths=0.5)

#3
r, rho = pearsonr(x=x, y=y)[0], spearmanr(a=x, b=y)[0]

#4
annotation = f'$r$ = {r:3.2f}, $\u03C1$ = {rho:3.2f}'
ax.text(0.05, 0.95, s=annotation, transform=ax.transAxes, fontsize=13,
        verticalalignment='top')

ax.set_xlabel('Sepal length [cm]')
ax.set_ylabel('Petal length [cm]')
ax.spines[['right', 'top']].set_visible(False)
ax.axis('equal')

title = 'Iris sepal length vs. petal length'
subtitle = ('Source: The Iris Dataset. UCI Machine Learning Repository '
            '(via Kaggle)')
set_title_and_subtitle(fig=fig, title=title, alignment='left',
                       subtitle=subtitle)

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