import itertools
import collections

import pycountry
import networkx as nx


_countries = {}


PREFIXES = set(['Republic of', ])
NAMES = {c.name.lower(): c for c in pycountry.countries.objects}


class CountryNotFound(Exception):
    def __init__(self, token, text):
        self.token = token
        self.text = text

    def __unicode__(self):
        return (u'Fuzzy search for "{}" returned no results (complete record: '
                '"{}")'.format(self.token, self.text))


_countries = {}


def get_country(text):
    tokens = text.split(', ')

    country_token = tokens[-1]

    if country_token in PREFIXES:
        country_token = tokens[-2] + ', ' + tokens[-1]

    if country_token not in _countries:
        country = None
        try:
            country = pycountry.countries.get(name=country_token)
        except KeyError:
            try:
                country = pycountry.countries.get(name=country_token)
            except KeyError:
                cl = country_token.lower()
                for k, country in NAMES.iteritems():
                    v = country.name.lower()
                    if cl in v or v in cl:
                        break
                    if country.alpha3 in country_token.split(' '):
                        break
                    if country.alpha2 in country_token.split(' '):
                        break
                else:
                    country = None
            if country:
                print 'Fuzzy search for "{}" matched "{}"'.format(
                    country_token, country.name)
            else:
                raise CountryNotFound(country_token, text)

        _countries[country_token] = country.name

    return _countries[country_token]


def get_countries(affiliations):
    countries = set()

    for a in affiliations:
        try:
            countries.add(get_country(a))
        except CountryNotFound as e:
            print unicode(e)

    return countries


def create(dataset):
    g = nx.Graph()

    affiliation_sets = dataset['author_affiliation']

    papers_count = collections.Counter()
    collaborations_count = collections.Counter()

    # Create graph
    for paper in affiliation_sets:
        if paper is None:
            continue
        countries = get_countries(paper.itervalues())
        g.add_nodes_from(countries)
        papers_count.update(countries)

        collaborations = list(itertools.combinations(countries, 2))
        collaborations_count.update(collaborations)
        g.add_edges_from(collaborations)

    # Set papers count
    for country, count in papers_count.iteritems():
        g.node[country]['papers'] = count

    # Set edge weight
    for (c1, c2), count in collaborations_count.iteritems():
        g[c1][c2]['weight'] = count

    return g

@@ -1,12 +1,23 @@
+import sys%0A
 import itert
@@ -81,16 +81,77 @@
  as nx%0A%0A
+from irco import models, logging%0A%0Alog = logging.get_logger()%0A
 %0A_countr
@@ -260,16 +260,75 @@
 bjects%7D%0A
+REPLACEMENTS = %7B%0A    'South Korea': 'Korea, Republic of'%0A%7D%0A
 %0A%0Aclass 
@@ -643,22 +643,57 @@
 country(
-text):
+institution):%0A    text = institution.name
 %0A    tok
@@ -740,32 +740,123 @@
 n = tokens%5B-1%5D%0A%0A
+    if country_token in REPLACEMENTS:%0A        country_token = REPLACEMENTS%5Bcountry_token%5D%0A%0A
     if country_t
@@ -1722,16 +1722,30 @@
    print
+ %3E%3Esys.stderr,
  'Fuzzy 
@@ -2016,28 +2016,27 @@
 untries(
-affi
+pub
 li
+c
 ation
-s
 ):%0A    c
@@ -2065,24 +2065,46 @@
 for 
-a in affi
+institution in pub
 li
+c
 ation
+.institution
 s:%0A 
@@ -2153,17 +2153,27 @@
 country(
-a
+institution
 ))%0A     
@@ -2221,16 +2221,30 @@
    print
+ %3E%3Esys.stderr,
  unicode
@@ -2282,23 +2282,23 @@
  create(
-dataset
+session
 ):%0A    g
@@ -2316,62 +2316,8 @@
 ()%0A%0A
-    affiliation_sets = dataset%5B'author_affiliation'%5D%0A%0A
     
@@ -2411,104 +2411,61 @@
     
-# Create graph%0A    for paper in affi
+for pub
 li
+c
 ation
-_sets:%0A        if paper is None:%0A            continue
+ in session.query(models.Publication):
 %0A   
@@ -2500,25 +2500,18 @@
 es(p
-aper.itervalues()
+ublication
 )%0A  

Update the country graph to use the database.