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

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

#2
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'

#3
axes = [left_ax, left_ax.twinx()]

#4
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]]
        ax.plot(x, y, marker='o')

#5
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)
#6
    ax.yaxis.set_major_locator(ticker.FixedLocator(ticks))
    ax.set_yticklabels(labels=labels, minor=False, fontweight="bold")

#7
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)

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.png', bbox_inches='tight', dpi=300)