import pandas as pd
import plotly.graph_objects as go
from colorbrewer import Set3 as palette

reporter_country = 'Germany'

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

#2
df.sort_values(by='Exports', inplace=True, ascending=False)
df['Exports'] = df['Exports']/10**6 

#3
cmap = palette[df.shape[0] + 1]
#4
node_colours = [f'rgb({c[0]:d},{c[1]:d},{c[2]:d})' for c in cmap]
#5
edge_colours = node_colours[1::]

#6
source_nodes_ids = [0] * df.shape[0]
target_nodes_ids = list(range(1, df.shape[0]+1))

#7
node_labels = [reporter_country] + [f'{pcountry} {w:4.1f}' for 
               pcountry, w in zip(df['Partner_country'], df['Exports'])]

#8
nodes = {'thickness':5, 'line':{'color':'black', 'width':0.5}, 
         'label':node_labels, 'color':node_colours}

#9
edges = {'source':source_nodes_ids, 'target':target_nodes_ids, 
         'value':df['Exports'], 'line':{'color':'black', 'width':0.2}, 
         'color':edge_colours, 'arrowlen':15}

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

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

#12
ptfig.add_annotation(
    {'text':(f'<b>{reporter_country} exports to the EU by partner '
             f'country (2021)</b>'),
     'x':0.0, 'y':1.3, 'xanchor':'left', 'showarrow':False,
     'font':{'color':'black', 'size':26}})

#13
ptfig.add_annotation(
    {'text':'Total exports, values are in bn USD',
     'x':0.0, 'y':1.2, 'xanchor':'left', 'showarrow':False,
     'font':{'color':'black', 'size':20}})
#14
ptfig.update_layout(font_family="Times New Roman")

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