from indra.statements import *


def has_location(stmt):
    if not stmt.context or not stmt.context.geo_location or \
            not stmt.context.geo_location.db_refs.get('GEOID'):
        return False
    return True


def has_time(stmt):
    if not stmt.context or not stmt.context.time:
        return False
    return True


def get_location(stmt):
    if not has_location(stmt):
        loc = None
    else:
        loc = stmt.context.geo_location.db_refs['GEOID']
    return loc


def get_time(stmt):
    if not has_time(stmt):
        time = None
    else:
        time = stmt.context.time
    return time


def location_matches(stmt):
    if isinstance(stmt, Event):
        context_key = get_location(stmt)
        matches_key = str((stmt.concept.matches_key(), context_key))
    elif isinstance(stmt, Influence):
        subj_context_key = get_location(stmt.subj)
        obj_context_key = get_location(stmt.obj)
        matches_key = str(stmt.matches_key(), subj_context_key,
                          obj_context_key)
    else:
        matches_key = stmt.matches_key()
    return matches_key


def event_location_refinement(st1, st2, hierarchies):
    ref = st1.refinement_of(st2, hierarchies)
    if not ref:
        return False
    if not has_location(st2):
        return True
    elif not has_location(st1):
        return False
    else:
        return st1.context.geo_location.db_refs['GEOID'] == \
               st2.context.geo_location.db_refs['GEOID']


def location_refinement(st1, st2, hierarchies):
    if type(st1) != type(st2):
        return False
    if isinstance(st1, Event):
        event_ref = event_location_refinement(st1, st2, hierarchies)
        return event_ref
    elif isinstance(st1, Influence):
        subj_ref = event_location_refinement(st1.subj, st2.subj,
                                             hierarchies)
        obj_ref = event_location_refinement(st1.obj, st2.obj,
                                            hierarchies)
        return subj_ref and obj_ref
    else:
        return st1.refinement_of(st2, hierarchies)


def event_location_time_matches(event):
    mk = location_matches(event)
    if not has_time(event):
        return mk
    time = get_time(event)
    matches_key = str((mk, time.start, time.end, time.duration))
    return matches_key


def location_time_matches(stmt):
    if isinstance(stmt, Event):
        return event_location_time_matches(stmt)
    elif isinstance(stmt, Influence):
        subj_mk = event_location_time_matches(stmt.subj)
        obj_mk = event_location_time_matches(stmt.obj)
        return str((stmt.matches_key(), subj_mk, obj_mk))
    else:
        return stmt.matches_key()


def event_location_time_refinement(st1, st2, hierarchies):
    ref = location_refinement(st1, st2, hierarchies)
    if not ref:
        return False
    if not has_time(st2):
        return True
    elif not has_time(st1):
        return False
    else:
        return st1.context.time.refinement_of(st2.context.time)


def location_time_refinement(st1, st2, hierarchies):
    if type(st1) != type(st2):
        return False
    if isinstance(st1, Event):
        return event_location_time_refinement(st1, st2, hierarchies)
    elif isinstance(st1, Influence):
        ref = st1.refinement_of(st2, hierarchies)
        if not ref:
            return False
        subj_ref = event_location_time_refinement(st1.subj, st2.subj,
                                                  hierarchies)
        obj_ref = event_location_time_refinement(st1.obj, st2.obj,
                                                 hierarchies)
        return subj_ref and obj_ref

@@ -42,32 +42,104 @@
 location(stmt):%0A
+    %22%22%22Return True if a Statement has grounded geo-location context.%22%22%22%0A
     if not stmt.
@@ -184,16 +184,16 @@
 on or %5C%0A
-
         
@@ -299,32 +299,87 @@
 has_time(stmt):%0A
+    %22%22%22Return True if a Statement has time context.%22%22%22%0A
     if not stmt.
@@ -467,32 +467,112 @@
 location(stmt):%0A
+    %22%22%22Return the grounded geo-location context associated with a Statement.%22%22%22%0A
     if not has_l
@@ -705,24 +705,87 @@
 time(stmt):%0A
+    %22%22%22Return the time context associated with a Statement.%22%22%22%0A
     if not h
@@ -900,32 +900,102 @@
 _matches(stmt):%0A
+    %22%22%22Return a matches_key which takes geo-location into account.%22%22%22%0A
     if isinstanc
@@ -1487,32 +1487,110 @@
 , hierarchies):%0A
+    %22%22%22Return True if there is a location-aware refinement between Events.%22%22%22%0A
     ref = st1.re
@@ -1930,32 +1930,109 @@
 , hierarchies):%0A
+    %22%22%22Return True if there is a location-aware refinement between stmts.%22%22%22%0A
     if type(st1)
@@ -2606,24 +2606,103 @@
 hes(event):%0A
+    %22%22%22Return Event matches key which takes location and time into account.%22%22%22%0A
     mk = loc
@@ -2914,24 +2914,97 @@
 ches(stmt):%0A
+    %22%22%22Return matches key which takes location and time into account.%22%22%22%0A
     if isins
@@ -3381,24 +3381,101 @@
 erarchies):%0A
+    %22%22%22Return True if there is a location/time refinement between Events.%22%22%22%0A
     ref = lo
@@ -3719,24 +3719,24 @@
 ext.time)%0A%0A%0A
-
 def location
@@ -3772,24 +3772,100 @@
 erarchies):%0A
+    %22%22%22Return True if there is a location/time refinement between stmts.%22%22%22%0A
     if type(

Add basic docstrings to custom assembly functions