"""
Amount of exports in millions of dollars per country in 2016, in logarithmic scale. 
Country shapes from http://www.naturalearthdata.com/downloads/110m-cultural-vectors/
"""

from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.colors import LogNorm
import numpy as np

from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from matplotlib.patches import PathPatch
from matplotlib.colors import LinearSegmentedColormap

#########################################################################################
# read values from data in wiki-table format
data = "".join(open("Export-wiki.txt").readlines())
data = data.replace("\n", "").replace("{","").replace("}","").replace(".","").replace(",",".").split("|-")[1:]
value={}
for line in data:
    line = line.split("|")
    try:
        val=float(line[3])
    except:
        continue
    key = line[2]
    value[key]=val
    
#########################################################################################
# draw empty worldmap
fig = plt.figure(figsize=(18,8.6))
plt.subplots_adjust(left=0.01, bottom=0.01, right=1.12, top=0.99)

m = Basemap(projection='robin',lon_0=0, llcrnrlat=-60,urcrnrlat=85, llcrnrlon=-180, urcrnrlon=180, resolution='l')
m.drawmapboundary()

#########################################################################################
# color country shapes
m.readshapefile('ne_110m_admin_0_countries/ne_110m_admin_0_countries', name='world', drawbounds=True, color='gray')

countries = []
undefined_countries = []
valueList  = []
for info, shape in zip(m.world_info, m.world):
    try:
        val = value[info["ADM0_A3"]]
    except:
        undefined_countries.append(Polygon(np.array(shape), True))
        continue

    countries.append(Polygon(np.array(shape), True))
    valueList.append(val)
    
valueArray = np.array(valueList)
print valueArray.min(), valueArray.max()

#########################################################################################
# colorbar, modified "gist_rainbow" theme
_gist_rainbow_data = (
        (0.000, (1.00, 1.00, 0.80)),
        #(0.030, (1.00, 0.00, 0.00)),
        (0.215, (1.00, 1.00, 0.00)),
        (0.400, (0.00, 1.00, 0.00)),
        (0.586, (0.00, 1.00, 1.00)),
        (0.770, (0.00, 0.00, 1.00)),
        (0.954, (1.00, 0.00, 1.00)),
        (1.000, (1.00, 0.00, 0.75)))
cm = colors.LinearSegmentedColormap.from_list("cm", _gist_rainbow_data, 256)
p = PatchCollection(countries, alpha=0.5,  zorder=3, cmap=cm, norm=LogNorm(vmin=200, vmax=2e6))
p.set_array(valueArray)

plt.gca().add_collection(p)
cb = fig.colorbar(p, shrink=0.6, pad = 0.02, drawedges=False)
cb.solids.set_edgecolor("face")

#########################################################################################
# set countries without data to lightgray
p2 = PatchCollection(undefined_countries, alpha=0.5,  zorder=3, cmap=LinearSegmentedColormap.from_list("lg", ["lightgray", "lightgray"]))
p2.set_array(np.ones((len(undefined_countries),)))
plt.gca().add_collection(p2)

#########################################################################################
# save and show
filename = "Export-Worldmap-2016"
plt.savefig(filename + ".svg")
plt.savefig(filename + "-1800px.png", dpi=100)
plt.savefig(filename + "-900px.png",  dpi=50)

plt.show()