#!/usr/bin/env python

# 3rd party modules
from lxml import etree
import mpu.pd
# Also need to install pandas


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

    countries = mpu.pd.countries
    countries = countries.set_index('en_name')
    country2dict = countries.to_dict('index')
    country2dict['Congo, Democratic Republic of'] = {'alpha-2': 'CD'}  # This is missing from the source dataset for some reason

    legend = [(mon_fri, '#B8D8D8'),  # gray-blue
              (mon_sat, 'orange'),
              (sun_thu, '#4F6367'),  # Dark gray
              (sat_thu, '#EB5D5D'),  # Red
              (sun_fri, '#60D838'),  # Green
              (mon_thu_sat, '#ff00ff'),  # Purple
              (mixed, '#00ffff')]  # Cyan

    colors = {}
    for collection, color in legend:
        for country in collection:
            if country in country2dict:
                code = country2dict[country]['alpha-2'].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()