import csv
import json

'''
Example structure
{
    "Canada": {
        "export-countries": {
            "China": 15800000000.0,
            "Japan": 8090000000.0,
            ...
        },
        "import-countries": {
            "China": 27300000000.0,
            "Germany": 10600000000.0,
            ...
        },
        ...

exports.csv
-----------
root,node,count
Canada,China,15800000000
Canada,Japan,8090000000

imports.csv
-----------
root,node,count
Canada,China,27300000000,
Canada,Germany,10600000000.0,

'''

G8 = set(["Canada", "China", "France", "Germany", "Italy", "Japan", "Russia", "United Kingdom", "United States"])

def create_links(origin_country, countries, result):
    for country in countries:
        if country in G8:
            value = int(countries[country])
            result.append([origin_country, country, value])

def process_data(data):

    exports = []
    imports = []

    for country in data:
        country_data = data[country]
        export_countries = country_data['export-countries']
        import_countries = country_data['import-countries']

        create_links(country, export_countries, exports)
        create_links(country, import_countries, imports)

    return {
        "exports" : exports,
        "imports" : imports
    }

def write_csv(filename, data):
    with open(filename, 'w') as csvfile:
        fieldnames = ['root', 'node', 'count']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        for row in data:
            writer.writerow({'root': row[0], 'node': row[1], 'count' : row[2]})

data_file = open('g8-trade.json', 'r')
data = json.load(data_file)

processed = process_data(data)

write_csv('exports.csv', processed['exports'])
write_csv('imports.csv', processed['imports'])