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

#1
df = pd.read_csv('data/top6-source-countries-us-imports.csv', comment='#')
start_yr, end_yr = df['Year'].min(), df['Year'].max()
ncountries = df['Country'].nunique()

#2
ncols = 2
nrows = np.ceil(ncountries/ncols).astype(int)  
grid_spec = {'hspace':0.5}
fig, ax = plt.subplots(nrows=nrows, ncols=ncols, gridspec_kw=grid_spec,
                       figsize=(6,7.0))

#3
ax = ax.flatten()

#4
ylims = (ymin := 0, ymax := 600)
yticks = np.linspace(ymin, ymax, num=4)

#5
for idx, (country, group) in enumerate(df.groupby(by='Country')):
        
#6
    x, y = group['Year'], group['Imports']
    ax[idx].plot(x, y, marker='o')
    
#7
    ax[idx].set_title(country)
    
#8
    ax[idx].set_xticks(df['Year'].unique()[0::2])
    ax[idx].set_xticks(df['Year'].unique()[1::2], minor=True)
    ax[idx].set_ylim(ylims)
    ax[idx].set_yticks(yticks)
    
    ax[idx].spines[['right', 'top', 'left']].set_visible(False)
    
title = (f'United States imports by source country '
         f'({start_yr}\u2012{end_yr})')
subtitle = ('Total imports in bn USD, goods only. '
            'Source: US Census Bureau.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                        alignment='left', h_offset=30)

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