# -*- coding: utf-8 -*-

from collections import namedtuple, Counter, OrderedDict
import re
import os

_ROOT = os.path.abspath(os.path.dirname(__file__))


def get_data_path(path):
    return os.path.join(_ROOT, 'data', path)


def read_table(filename, usecols=(0, 1), sep='\t', comment='#', encoding='utf-8', skip=0):
    """Parse data files from the data directory

    Parameters
    ----------
    filename: string
        Full path to file

    usecols: list, default [0, 1]
        A list of two elements representing the columns to be parsed into a dictionary.
        The first element will be used as keys and the second as values. Defaults to
        the first two columns of `filename`.

    sep : string, default '\t'
        Field delimiter.

    comment : str, default '#'
        Indicates remainder of line should not be parsed. If found at the beginning of a line,
        the line will be ignored altogether. This parameter must be a single character.

    encoding : string, default 'utf-8'
        Encoding to use for UTF when reading/writing (ex. `utf-8`)

    skip: int, default 0
        Number of lines to skip at the beginning of the file

    Returns
    -------
    A dictionary with the same length as the number of lines in `filename`
    """

    with open(filename, 'r') as f:
        # skip initial lines
        for _ in range(skip):
            next(f)

        # filter comment lines
        lines = (line for line in f if not line.startswith(comment))

        d = dict()
        for line in lines:
            columns = line.split(sep)
            key = columns[usecols[0]].lower()
            value = columns[usecols[1]].rstrip('\n')
            d[key] = value
    return d


def build_index():
    """Load information from the data directory

    Returns
    -------
    A namedtuple with three fields: nationalities cities countries
    """

    nationalities = read_table(get_data_path('nationalities.txt'), sep=':')

    # parse http://download.geonames.org/export/dump/countryInfo.txt
    countries = read_table(
        get_data_path('countryInfo.txt'), usecols=[4, 0], skip=1)

    # parse http://download.geonames.org/export/dump/cities15000.zip
    cities = read_table(get_data_path('cities15000.txt'), usecols=[1, 8])

    # load and apply city patches
    city_patches = read_table(get_data_path('citypatches.txt'))
    cities.update(city_patches)

    Index = namedtuple('Index', 'nationalities cities countries')
    return Index(nationalities, cities, countries)


class GeoText(object):

    """Extract cities and countries from a text

    Examples
    --------

    >>> places = GeoText("London is a great city")
    >>> places.cities
    "London"

    >>> GeoText('New York, Texas, and also China').country_mentions
    OrderedDict([(u'US', 2), (u'CN', 1)])

    """

    index = build_index()

    def __init__(self, text):
        city_regex = r"[A-Z]+[a-z]*(?:[ '-][A-Z]+[a-z]*)*"
        candidates = re.findall(city_regex, text)
        self.countries = [each for each in candidates
                          if each.lower() in self.index.countries]
        self.cities = [each for each in candidates
                       if each.lower() in self.index.cities
                       # country names are not considered cities
                       and each.lower() not in self.index.countries]
        self.nationalities = [each for each in candidates
                              if each.lower() in self.index.nationalities]

        # Calculate number of country mentions
        self.country_mentions = [self.index.countries[country.lower()]
                                 for country in self.countries]
        self.country_mentions.extend([self.index.cities[city.lower()]
                                      for city in self.cities])
        self.country_mentions.extend([self.index.nationalities[nationality.lower()]
                                      for nationality in self.nationalities])
        self.country_mentions = OrderedDict(
            Counter(self.country_mentions).most_common())

if __name__ == '__main__':
    print(GeoText('In a filing with the Hong Kong bourse, the Chinese cement producer said ...').countries)

@@ -2907,16 +2907,19 @@
 -Z%5D+%5Ba-z
+%C3%A0-%C3%BA
 %5D*(?:%5B '
@@ -2930,16 +2930,19 @@
 -Z%5D+%5Ba-z
+%C3%A0-%C3%BA
 %5D*)*%22%0A  

Adding ability to recognize more cities