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

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

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

colours, styles = ['#66c2a5', '#fc8d62', '#8da0cb'], ['o', 's', 'D']

#1
img_folder, ext = 'resources/pictures', 'png' 
names = ['Iris-setosa', 'Iris-versicolor', 'Iris-virginica']

sns.scatterplot(data=df, x='SepalLengthCm', y='PetalWidthCm',
                hue='Species', style='Species', palette=colours,
                markers=styles, hue_order=names, style_order=names, 
                edgecolor='k', linewidth=0.5, legend=False)

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

#2
x_bounds = ax.get_xbound() 
y_bounds = ax.get_ybound()
ax_width = x_bounds[1] - x_bounds[0]
ax_height = y_bounds[1] - y_bounds[0]
ax_form_factor = ax_height/ax_width

#3
h_starting_pos, v_starting_pos = 0.45, 0.10 
h_spacing, v_spacing = 0.05, 0.075/ax_form_factor
width, height = 0.15, 0.15/ax_form_factor

#4
for idx, item in enumerate(zip(colours, names)):
    
#5
    h_pos = h_starting_pos + idx*(h_spacing + width)
    v_pos = v_starting_pos + idx*v_spacing
    
#6
    axin = ax.inset_axes(bounds=[h_pos,v_pos,width,height])
    image = plt.imread(f'{img_folder}/{item[1]}.{ext}', format=ext)
    axin.imshow(image) 
    
#7
    axin.set_xticks([])
    axin.set_yticks([])
    axin.set_xlabel(f'{item[1]}', color=item[0])
    
#8
    axin.spines[['left','bottom','right','top']].set_linewidth(3)
    axin.spines[['left','bottom','right','top']].set_color(item[0])
   
title = 'Iris dataset: Sepal length vs. petal width'
subtitle = 'Source: The Iris Dataset. UCI Machine Learning Repository (via Kaggle)'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle, 
                       alignment='left', h_offset=15, v_offset=15)

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