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


fig, left_ax = plt.subplots(nrows=1, ncols=1, figsize=(4,5))

df = pd.read_csv('data/top-purchasers-us-exports-2015-2021.csv',
                 comment='#')
start_end_yr = (start_yr := 2015, end_yr := 2021)
var = 'Perc_total_exp'

axes = [left_ax, left_ax.twinx()]

#1
cmap = {'Increasing':'#fdae61', 'Decreasing':'#2b83ba'}
countries = df.groupby(by=['Country'])
for _, country in countries:
    for ax in axes:
        x = [start_yr, end_yr]
        y = [country[country['Year'] == yr][var].values for 
             yr in [start_yr, end_yr]]
#2
        trend = 'Increasing' if (y[1] > y[0]) else 'Decreasing' 
        ax.plot(x, y, marker='o', color=cmap[trend])
        
def generate_label(x): return f"{x['Country']} ({x[var]}%)"
for idx, ax in enumerate(axes):
    df_slice = df[df['Year'] == start_end_yr[idx]]    
    ticks = df_slice[var]
    labels = df_slice.agg(generate_label, axis=1)
    ax.yaxis.set_major_locator(ticker.FixedLocator(ticks))
    ax.set_yticklabels(labels=labels, minor=False, fontweight="bold")
        
x_extra_space = 0.25
xlims = start_yr-x_extra_space, end_yr+x_extra_space  
for ax in axes:
    ax.set_xlim(left=xlims[0], right=xlims[1])   
    ax.xaxis.set_major_locator(ticker.FixedLocator(start_end_yr))
    ax.spines[['bottom','top']].set_visible(False)

#3
legend_entries = list()
for key, value in cmap.items():
    legend_entries.append(mpatches.Patch(color=value, label=key))        

#4
ax.legend(handles=legend_entries, loc='upper center', ncol=2, 
          bbox_to_anchor=(0.5,-0.05), title='Overall trend', 
          frameon=False)

title = (f'Top purchasers of United States exports: '
         f'{start_yr} vs {end_yr}')
subtitle = (f'Percentage of total exports in goods. '
            f'Source: US Census Bureau.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle, 
                       alignment='left', h_offset=-55, v_offset=0)

fig.savefig('charts/slope-chart-legend.png', bbox_inches='tight',
            dpi=300)