import matplotlib.pyplot as plt
import numpy as np
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=(6,6))

#1
df_temp_data = pd.read_csv(
    'data/uk-metoff-temperature-data-2013-2022.csv', comment='#')
df_station_data = pd.read_csv(
    'data/uk-metoff-station-data.csv', comment='#')

#2
df_station_data.sort_values(by='Latitude', ascending=False, inplace=True)
sel_stations = df_station_data['Station name'][0::3]
df = df_temp_data.loc[df_temp_data['Station name'].isin(sel_stations)]

#3
month = ('July', 7); df = df[df['mm'] == month[1]]
#4
df = df.pivot(index='Station name', columns='Year', values='tmax')
#5
df = df.reindex(labels=sel_stations) 

#6
sns.heatmap(data=df, ax=ax, linewidths=1.5, linecolor='lightgray', 
            cmap='Reds') 

#7
ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
ax.set_yticklabels(ax.get_yticklabels(), rotation=0)

#8
ax.set_xlabel(None), ax.set_ylabel(None) 

start_yr, end_yr = min(df_temp_data['Year']), max(df_temp_data['Year'])
title = (f'UK temperature data at selected historic stations '
         f'({start_yr}-{end_yr})')
subtitle = (f'Mean daily low in {month[0]} [\u00B0C].'
            f' Source: UK Met Office.')
set_title_and_subtitle(title=title, subtitle=subtitle, fig=fig, 
                       alignment='left', h_offset=-24)

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