from itertools import product
import pandas as pd
import plotly.graph_objects as go
from colorbrewer import Accent as palette

df = pd.read_csv('data/eu3-to-sa4-trade-2021.csv', comment='#')
df['Exports'] = df['Exports']/10**6 

#1
reporter_countries = set(df['Reporter_country'])
partner_countries = set(df['Partner_country'])

#2
all_countries = list(reporter_countries) + list(partner_countries)
country_ids = {country: i for i, country in enumerate(all_countries)}

#3
labels = list()
for country in all_countries:
    value = df[(df['Reporter_country'] == country) | 
               (df['Partner_country'] == country)]['Exports'].sum()
    labels.append(f'{country} ({value:4.1f})')

#4
alpha = 0.7

#5
cmap = palette[len(reporter_countries)]
colours_reporter = [f'rgba({c[0]:d},{c[1]:d},{c[2]:d},{alpha:3.2f})' 
                    for c in cmap]
#6
colours_partner = ['rgba(127,127,127,0.0)'] * len(partner_countries)
all_colours = colours_reporter + colours_partner

source_nodes_ids, target_nodes_ids, weights, edge_colours =\
 [list() for i in range(4)]

#7
for reporter_country, partner_country in product(reporter_countries,
                                                 partner_countries):
    
#8
    source_nodes_ids.append(country_ids[reporter_country])
    target_nodes_ids.append(country_ids[partner_country])
    
#9
    df_slice = df[((df['Reporter_country'] == reporter_country) & 
                   (df['Partner_country'] == partner_country))]
    weights.append(df_slice['Exports'].values[0])
    
#10
    edge_colours.append(all_colours[country_ids[reporter_country]])
    
nodes = {'thickness':5, 'line':{'color':'black', 'width':0.5}, 
         'label':labels, 
         'color':colours_reporter + colours_partner}

edges = {'source':source_nodes_ids, 'target':target_nodes_ids, 
         'value':weights, 'line':{'color':'black', 'width':0.2}, 
         'color':edge_colours, 'arrowlen':15}

ptfig = go.Figure(data=[go.Sankey(node=nodes, link=edges)])

ptfig.update_layout(font_family='Times New Roman', font_color='black', 
                    font_size=16)

ptfig.add_annotation(
    {'text':'<b>Exports from EU to Latin America (2021)</b>',
     'x':0.0, 'y':1.3, 'xanchor':'left', 'showarrow':False,
     'font':{'color':'black', 'size':26}})
ptfig.add_annotation(
    {'text':'Total exports [bn USD]. Source: The World Bank/WITS.',
     'x':0.0, 'y':1.2, 'xanchor':'left', 'showarrow':False,
     'font':{'color':'black', 'size':20}})

ptfig.update_layout(
    font_family="Times New Roman",
)

ptfig.write_image(file='charts/sankey-many-to-many.png', scale=3)