"""
Number of deaths of infants under five years per 1,000 live births for 2016. 
Data from https://data.worldbank.org/indicator/SH.DYN.MORT?end=2016&start=2016&view=map
Country shapes from http://www.naturalearthdata.com/downloads/110m-cultural-vectors/
Created with Python and Matplotlib Basemap Toolkit.
"""

from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
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 mortality data
data = open("API_SH.DYN.MORT_DS2_en_csv_v2/API_SH.DYN.MORT_DS2_en_csv_v2.csv").readlines()
mortality={}
for line in data[5:]:
    line = line.replace('",','|').replace(",","").replace('"',"").strip().split("|")
    if len(line)<10:
        continue
    try:
        value=float(line[-2])
    except:
        continue
    key = line[1]
    mortality[key]=value

#########################################################################################
# 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 = []
mortList  = []
for info, shape in zip(m.world_info, m.world):
    try:
        mort = mortality[info["ADM0_A3"]]
    except:
        undefined_countries.append(Polygon(np.array(shape), True))
        continue

    countries.append(Polygon(np.array(shape), True))
    mortList.append(mort)
    
mortArray = np.array(mortList)
print mortArray.min(), mortArray.max()
ticks = np.linspace(0, 130,14)

#########################################################################################
# colorbar
cm = LinearSegmentedColormap.from_list("cm", ["#81fcff", "#19ff19", "#e2f000", "#ffaa31", "#ff8e51", "#ff6969", "#8c0000", "#3c0000"])
p = PatchCollection(countries, alpha=0.5,  zorder=3, cmap=cm)
p.set_array(mortArray)
p.set_clim([ticks.min(), ticks.max()])

plt.gca().add_collection(p)
cb = fig.colorbar(p, ticks = ticks, shrink=0.6, pad = 0.02)

#########################################################################################
# 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 = "infant-under-five-mortality-worldmap-2016"
plt.savefig(filename + ".svg")
plt.savefig(filename + "-1800px.png", dpi=100)
plt.savefig(filename + "-900px.png",  dpi=50)

plt.show()