from datatool import datatool
import numpy as np

##############################################################
# Write your code in the functions below

def find_first_import_row(D):
    # COPY your code from an earlier exercise

def country_list(D, column_name, N):
    # Sort by column name, then get the countries in the top N rows
    dt.sort_by_multiple_columns([column_name], [False])
    D = dt.get_data_as_array()
    countries = []
    # WRITE YOUR CODE HERE to append to the countries list those
    # countries in the first N rows of D

    return countries

def list_diff(L1, L2):
    # Create a new list of elements in L from L1 that are NOT in L2
    L = []
    # WRITE YOUR CODE HERE

    return L

# We've written this for you. Read to see how the gravity model is used.
def make_gravity_list(D):
    gravity_list = []
    for i in range(D.shape[0]):
        gravity_list.append(0.0)
    first_import_row = find_first_import_row(D)
    for i in range(first_import_row):
        for j in range(first_import_row, D.shape[0]):
            if D[j,1] == D[i,1]:       # same country
                # GDP per person:
                US_gdp_pop = 16598099 / 312355
                other_gdp_pop = D[i,5] / D[i,13]
                # The GDP product:
                gdp_prod_scaled = US_gdp_pop * other_gdp_pop 
                # Calculate the ratio:
                value = (D[i,4] + D[j,4]) / gdp_prod_scaled
                gravity_list[i] = value
                gravity_list[j] = value
                break
    return gravity_list

##############################################################
# The program starts here

dt = datatool()

dt.load_csv('US_trade_2014.csv')

dt.sort_by_multiple_columns(['trade_direction'], [True])

D = dt.get_data_as_array()

# Compute the gravity for each country against the U.S.
gravity_list = make_gravity_list(D)

# This is how to add a column to the existing data:
dt.add_column('gravity', gravity_list)

# We'll limit the data to only the export rows:
first_import_row = find_first_import_row(D)
dt.use_n_rows(first_import_row)

# Extract the high-gdp countries into a list:
high_gdp = country_list(D, 'gdp', 20)

# Extract the high-gravity countries into a list:
high_gravity = country_list(D, 'gravity', 20)

# Print these:
print('High gdp', high_gdp)
# High gdp ['China', 'India', 'Japan', 'Germany', 'Russia', 'Brazil', 'France', 'United Kingdom', 'Indonesia', 'Italy', 'Mexico', 'Canada', 'Turkey', 'Spain', 'Saudi Arabia', 'Iran', 'Australia', 'Nigeria', 'Poland', 'Thailand']

print('High_gravity', high_gravity)
# High_gravity ['China', 'Mexico', 'Canada', 'India', 'Vietnam', 'Japan', 'Brazil', 'Germany', 'Philippines', 'Indonesia', 'Thailand', 'Colombia', 'United Kingdom', 'Honduras', 'Venezuela', 'Bangladesh', 'France', 'Malaysia', 'Italy', 'Nigeria']

# Which countries are in the high gdp that are not in high-gravity?
gdp_minus_gravity = list_diff(high_gdp, high_gravity)

print('High gdp not in high gravity', gdp_minus_gravity)
# High gdp not in high gravity ['Russia', 'Turkey', 'Spain', 'Saudi Arabia', 'Iran', 'Australia', 'Poland']