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 ='#')
df['Total_trade'] = df['Exports'] + df['Imports'] 

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

flow = np.zeros((len(nodes),len(nodes)))
pairs = permutations(nodes, r=2)
for pair in pairs:

    idx_reporter_country = nodes.index(pair[0])
    idx_partner_country = nodes.index(pair[1])
    
    tot_trade = df[(df['Reporter_country'] == pair[0]) &
                   (df['Partner_country'] == pair[1])]['Total_trade']
#2
    flow[idx_reporter_country][idx_partner_country] = tot_trade

#3
chord_diagram(mat=flow, names=nodes, ax=ax, directed=False, cmap='Set1')

title = f'Total trade between 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-undirected.png', dpi=300, 
            bbox_inches='tight')