import ujson as json

from sift.corpora import wikicorpus
from sift.dataset import ModelBuilder, Model, Relations

from sift import logging
log = logging.getLogger()

ENTITY_PREFIX = 'Q'
PREDICATE_PREFIX = 'P'

class WikidataCorpus(ModelBuilder, Model):
    @staticmethod
    def iter_item_for_line(line):
        line = line.strip()
        if line != '[' and line != ']':
            yield json.loads(line.rstrip(',\n'))

    def build(self, sc, path):
        return sc\
            .textFile(path)\
            .flatMap(self.iter_item_for_line)\
            .map(lambda i: (i['id'], i))

    @staticmethod
    def format_item((wid, item)):
        return {
            '_id': wid,
            'data': item
        }

class WikidataRelations(ModelBuilder, Relations):
    """ Prepare a corpus of relations from wikidata """
    @staticmethod
    def iter_relations_for_item(item):
        for pid, statements in item.get('claims', {}).iteritems():
            for statement in statements:
                datatype = statement['mainsnak'].get('datatype', None)
                if datatype == 'wikibase-item':
                    tid = statement['mainsnak']['datavalue']['value']['numeric-id']
                    yield pid, int(tid)
                elif datatype == 'url':
                    yield pid, statement['mainsnak']['datavalue']['value']
                elif datatype == 'time':
                    yield pid, statement['mainsnak']['datavalue']['value']['time']

    def build(self, corpus):
        entities = corpus\
            .filter(lambda item: item['_id'].startswith(ENTITY_PREFIX))

        entity_labels = entities\
            .map(lambda item: (item['_id'], item['data'].get('labels', {}).get('en', {}).get('value', None)))\
            .filter(lambda (pid, label): label)\
            .map(lambda (pid, label): (int(pid[1:]), label))

        wiki_entities = entities\
            .map(lambda item: (item['data'].get('sitelinks', {}).get('enwiki', {}).get('title', None), item['data']))\
            .filter(lambda (e, _): e)\
            .cache()
       
        predicate_labels = corpus\
            .filter(lambda item: item['_id'].startswith(PREDICATE_PREFIX))\
            .map(lambda item: (item['_id'], item['data'].get('labels', {}).get('en', {}).get('value', None)))\
            .filter(lambda (pid, label): label)\
            .cache()

        relations = wiki_entities\
            .flatMap(lambda (eid, item): ((pid, (value, eid)) for pid, value in self.iter_relations_for_item(item)))\
            .join(predicate_labels)\
            .map(lambda (pid, ((value, eid), label)): (value, (label, eid)))

        return relations\
            .leftOuterJoin(entity_labels)\
            .map(lambda (value, ((label, eid), value_label)): (eid, (label, value_label or value)))\
            .groupByKey()\
            .mapValues(dict)

@@ -985,16 +985,89 @@
 ements:%0A
+                if statement%5B'mainsnak'%5D.get('snaktype') == 'value':%0A    
         
@@ -1125,16 +1125,12 @@
 ype'
-, None
 )%0A
+  
     
@@ -1141,16 +1141,18 @@
         
+  
 if datat
@@ -1175,16 +1175,20 @@
 -item':%0A
+    
         
@@ -1275,32 +1275,36 @@
                 
+    
 yield pid, int(t
@@ -1315,32 +1315,36 @@
                 
+   
+ 
 elif datatype ==
@@ -1351,16 +1351,20 @@
  'url':%0A
+    
         
@@ -1438,32 +1438,36 @@
                 
+    
 elif datatype ==
@@ -1471,24 +1471,28 @@
  == 'time':%0A
+    
             

Handle relations with no value in recent wikidata dumps