#!/usr/bin/env python

# 3rd party modules
from lxml import etree
import mpu.pd


def main(filename="weekend.svg"):
    mon_fri = ['Dominican Republic', 'Lebanon', 'Serbia', 'Albania', 'Angola',
               'Argentina', 'Armenia', 'Azerbaijan', 'Austria', 'Australia',
               'Benin',
               'Belarus', 'Belgium', 'Brazil', 'Burundi', 'Bulgaria', 'Canada',
               'Cambodia', 'Cameroon', 'Chile', 'China', 'Croatia',
               'Costa Rica', 'Czech Republic', 'Denmark', 'Ethiopia',
               'Estonia', 'Finland', 'France', 'Gabon',
               'Gambia', 'Germany', 'Ghana', 'Greece', 'Hungary', 'Indonesia',
               'Ireland',
               'Italy', 'Côte d\'Ivoire', 'Japan', 'Kazakhstan', 'Kenya',
               'Lao People\'s Democratic Republic', 'Latvia', 'Lesotho',
               'Lithuania', 'Madagascar',
               'Malawi', 'Mali', 'Malta', 'Mauritania', 'Mongolia', 'Morocco',
               'Mozambique', 'Netherlands', 'New Zealand', 'Nigeria',
               'Norway',
               'Pakistan', 'Poland', 'Portugal', 'Romania', 'Russia', 'Rwanda',
               'Senegal',
               'Singapore', 'Slovakia', 'Spain', 'Sri Lanka',
               'South Africa', 'South Korea', 'Suriname', 'Swaziland',
               'Sweden', 'Switzerland',
               'Seychelles', 'Taiwan', 'Tanzania', 'Togo', 'Thailand',
               'Trinidad and Tobago', 'Tunisia', 'Turkey', 'Ukraine',
               'United Kingdom', 'United States',
               'Vietnam', 'Congo (Democratic Republic of the)', 'Zambia',
               'Zimbabwe']
    mon_sat = ['Equatorial Guinea', 'Hong Kong', 'Mexico', 'Philippines',
               'Uganda', 'India', 'Colombia']
    sun_thu = ['Afghanistan', 'Algeria', 'Bahrain', 'Egypt', 'Iraq', 'Israel',
               'Jordan', 'Kuwait', 'Libya', 'Maldives', 'Oman', 'Qatar',
               'Saudi Arabia', 'Sudan', 'Syria', 'Yemen',
               'United Arab Emirates', 'Bangladesh']
    sat_thu = ['Djibouti', 'Palestine', 'Iran']
    sun_fri = ['Nepal']
    mon_thu_sat = ['Brunei Darussalam']
    mixed = ['Malaysia']

    countries = mpu.pd.countries
    countries = countries.set_index('English short name (upper/lower case)')
    country2dict = countries.to_dict('index')

    legend = [(mon_fri, '#B8D8D8'),
              (mon_sat, 'orange'),
              (sun_thu, '#4F6367'),
              (sat_thu, 'red'),
              (sun_fri, 'tomato'),
              (mon_thu_sat, '#ff00ff'),
              (mixed, '#00ffff')]

    colors = {}
    for collection, color in legend:
        for country in collection:
            if country in country2dict:
                code = country2dict[country]['Alpha-2 code'].lower()
                colors[code] = color
            else:
                print('Could not find the Alpha-2 code for \'{}\' in mpu'
                      .format(country))

    tree = read_svg(filename)
    root = tree.getroot()
    for country_id in colors.keys():
        element = root.xpath('//*[@id="{}"]'.format(country_id))
        if len(element) == 1:
            element = element[0]
            element.set("style", 'fill: {};'.format(colors[country_id]))
            for el in element.getchildren():
                el.set("style", 'fill: {};'.format(colors[country_id]))
        else:
            print("Wrong length for id={}".format(country_id))
    write_svg(root)


def read_svg(filename):
    """
    Read an SVG file.

    Parameters
    ----------
    filename : str

    Returns
    -------
    tree : lxml.ElementTree
    """
    with open(filename, 'r') as f:
        tree = etree.parse(f)
    return tree


def write_svg(root, filename='weekend-auto.svg'):
    """
    Read an SVG file.

    Parameters
    ----------
    root : lxml.Element
    filename : str
    """
    etree.ElementTree(root).write(filename, pretty_print=True)


if __name__ == '__main__':
    main()