#396 Seaborn pairplot example with lowess regression

import matplotlib.pyplot as plt
import seaborn as sns

def plot_extra(x, y, **kwargs):
     if kwargs['label'] == first_label:
          sns.regplot(data=kwargs['data'], x=x.name, y=y.name, lowess=True, scatter=False, color=kwargs['color'])

df = sns.load_dataset('iris')
first_label = df['species'][0]
pg = sns.pairplot(df, hue='species', plot_kws={'alpha': 0.8}, palette='rainbow')
pg.map_offdiag(plot_extra, color='crimson', data=df)
legend_dict = {h.get_label(): h for h in pg.legend.legend_handles}  # the existing legend items
legend_dict['lowess regression'] = pg.axes[0, 1].lines[0]
# add the first line object of a non-diagonal ax to the legend
pg.legend.remove()  # remove existing legend
pg.add_legend(legend_dict, label_order=legend_dict.keys(), title='')  # create the new legend
plt.savefig("ex396.png", dpi=100)
plt.show()