from itertools import permutations
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from mpl_chord_diagram import chord_diagram
from mpl_ornaments.titles import set_title_and_subtitle

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5,6))

#1
df = pd.read_csv('data/eu7-trade-2021.csv', comment='#')

#2
reporter_countries = set(df['Reporter_country'])
partner_countries = set(df['Partner_country'])
nodes = list(reporter_countries.intersection(partner_countries))
nodes.sort()

#3
flow = np.zeros((len(nodes),len(nodes)))
pairs = permutations(nodes, r=2)
for pair in pairs:
#4
    idx_reporter_country = nodes.index(pair[0])
    idx_partner_country = nodes.index(pair[1])
    
    exports = df[(df['Reporter_country'] == pair[0]) &
                 (df['Partner_country'] == pair[1])]['Exports']
#5
    flow[idx_reporter_country][idx_partner_country] = exports

#6
chord_diagram(mat=flow, names=nodes, ax=ax, directed=True, cmap='Set1')

title = f'Exports from/to the seven largest EU economies (2021)'
subtitle = 'Source: The World Bank/WITS.'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle, 
                       alignment='left')

fig.savefig('charts/chord-diagram-directed.png', dpi=300, 
            bbox_inches='tight')