%tensorflow_version 2.x
!test -d language_repo || git clone https://github.com/google-research/language.git language_repo
%cd /content/language_repo/language/nql
!pwd
!pip install .
%cd
!pwd

import math
import time
import importlib

import numpy as np
import tensorflow as tf
import nql
from nql import util
  
# Prints a dictionary or list in a sorted way.
def printd(obj, top=None):
  if isinstance(obj, dict):
    print(sorted(obj.items())[:top])
  elif isinstance(obj, list):
    print(sorted(obj)[:top])
  else:
    print(obj)

"""
Explanation: Getting started

Imports
Defining the KG schema
Loading KG triples

Following relations in the KG
Conditionals
Quantifying over relations
Under the hood - NQL and TensorFlow

Weighted sets and counting paths
Why weights are useful
Following relations via sparse matrix multiplication


Getting started
Imports
NQL is a PathQuery-like language which is closely integrated with Tensorflow.  This lab walks you through some simple use cases for NQL, the Neural Query Language.  The focus is on explaining the semantics of the query language - another CoLab looks at using NQL in conjunction with Tensorflow for learning.
End of explanation
"""


context = nql.NeuralQueryContext()
relations = [
  'aunt', 'brother', 'daughter', 'father', 'husband', 'mother', 
  'nephew', 'niece', 'sister', 'son', 'uncle', 'wife' 
]
for r in relations:
  context.declare_relation(r, 'person_t', 'person_t')

print('declared relations', sorted(context.get_relation_names()))
print('declared types', sorted(context.get_type_names()))


"""
Explanation: Defining the KG schema
Now, let's set up a small KG to run over.  The one we'll work with is inspired by a very old Hinton paper on learning relations. It has one type - a person type - and twelve relations, which correspond to familial relations. 
We start out by setting up a context object, to hold the data, and defining a schema.  By convention a type name will end with "_t". When we declare a relation, we will provide the name of the type of its first and second arguments.  So to define the schema for this KG we will do this:
End of explanation
"""


with open('/content/language_repo/language/nql/demos/data/royal92/royal_family.tsv') as f:
  context.load_kg(files=f)

"""
Explanation: Loading KG triples
Now let's load in some data.  These are the relationships between members of several European royal families.  Loading these will define the relations that we declared.
End of explanation
"""


with open('/content/language_repo/language/nql/demos/data/royal92/royal_family.tsv') as f:
  print(f.readlines()[:3])

"""
Explanation: The original input is stored as tab-separated triples, each triple being a relation name, followed by a person name.  For instance the first few lines are as follows.
End of explanation
"""


queen_vic = context.one('Victoria of house of Hanover', 'person_t')
henry8 = context.one('Henry_VIII of house of Tudor', 'person_t')

print(queen_vic.eval())
print(henry8.eval())

"""
Explanation: Following relations in the KG
Now that a KG is loaded in, we can use the context object as a factory to create some NQL expressions.  The one method creates a singleton set containing one entity, which is specified via its name and type.
End of explanation
"""


printd(henry8.wife().eval())

"""
Explanation: Expressions in NQL always evaluate to sets of entities.  So, if you want to inspect a particular entity, then you need to construct a singleton set containing that entity.    For example, the expressions below find people in the KG (i.e., entities of type 'person_t').  Like Tensorflow expressions NQL expressions are lazy, and you should evaluate them in the context of a session.  
The default output of nql_expression.eval() is a Python dictionary mapping entity names to floats.  The floats make it possible to have a "soft" or weighted set of entities.  We'll come back to this feature later, but for now we can ignore the floats.
We can also follow relations in the KG, for example, asking 'who were the wives of Henry VIII?'
End of explanation
"""


printd(queen_vic.follow('husband').eval())
printd(queen_vic.follow('wife', -1).eval())
printd(queen_vic.wife(-1).eval())
printd(queen_vic.son().son().eval())

"""
Explanation: Alternatively you can follow relations by name, or follow the inverse of a relation.  You can also chain relations.
End of explanation
"""


print('queen victoria\'s children')
print(' - sons:')
printd(queen_vic.son().eval())

print(' - daughters:')
printd(queen_vic.daughter().eval())

print(' - all:')
printd((queen_vic.son() | queen_vic.daughter()).eval())

printd((henry8.wife() & context.one('Ferdinand_V', 'person_t').daughter()).eval())

"""
Explanation: You can take unions and intersections of sets.
End of explanation
"""


def child(x): return x.son() | x.daughter()
def parent(x): return x.mother() | x.father()
def grandfather(x): return parent(x).father()

printd(child(queen_vic).eval())
printd(grandfather(henry8).eval())

"""
Explanation: (By the  way, * and + are synonyms for & and |, so you can use either notation).
And as in TensorFlow, you can use Python functions which return expressions as macros, to build up more complex relationships.
End of explanation
"""


def male(x): return (x.mother().son() & x)
def female(x): return (x.mother().daughter() & x)

printd(male(context.all('person_t')).eval(), 5)
printd(female(context.all('person_t')).eval(), 5)



"""
Explanation: Also notice that the newly-defined constructs are functions: e.g., male(S) returns a new set S' which contains the male members of S.  If it seems awkward to switch from "chaining-style" syntax (e.g., x.son().son()) to "function-call" syntax, there is a way around that, which we will discuss below.
Conditionals
This dataset has no sex or gender indications, so here's a way to identify the men and women in the dataset.  Note the use of the new construct context.all(type) which returns a set of all entities in a type.
End of explanation
"""


def opposite_sex_parent(x):
  return x.father().if_any(female(x)) | x.mother().if_any(male(x))

print('opposite_sex_parent of henry:')
printd(opposite_sex_parent(henry8).eval())
print('opposite_sex_parent of victoria:')
printd(opposite_sex_parent(queen_vic).eval())

"""
Explanation: Note that the weights in these sets are not uniform -- we'll get back to this below, in the "under the hood" section.
Now, however, we will use male and female to demonstrate the syntax for conditional relation-following.  As a somewhat artificial example, suppose we want to define opposite_sex_parent as the father of a female person, and mother of a male person.  We could do this as follows.
End of explanation
"""


nuclear_family_relations = ['son', 'daughter', 'father', 'mother', 'brother', 'sister', 'wife', 'husband']
g = context.construct_relation_group('nuclear_family_g', 'person_t', 'person_t', nuclear_family_relations)
printd(g.__dict__)

"""
Explanation: Quantifying over relations
NQL lets you "follow" specific relations.  Sometimes it's useful to variabilize the relations you follow.  We already saw the syntax x.follow(r) where r is a Python variable bound to a string naming a relation.  Here we'll discuss a way to push these "relation variables" down into Tensorflow.
We will first construct a "relation group" consisting of a selected subset of relations.  By convention relation group names end with "_g".  You usually won't need the return value of construct_relation_group but it's printed to help explain what's happens under the hood, below.
End of explanation
"""


print('all relations:')
printd(context.get_relation_names())
print('all types:')
printd(context.get_type_names())
print('new relations:')
printd(list(filter(lambda x:x.startswith(g.name), context.get_relation_names())))
print('new types:')
printd(list(filter(lambda x:x.startswith(g.name), context.get_type_names())))

"""
Explanation: The effect of this call is to add two new types and four new relations to the KG.
End of explanation
"""


for rel_id in range(context.get_max_id('nuclear_family_g')):
  print(context.get_entity_name(rel_id, 'nuclear_family_g'))

"""
Explanation: The instances of nuclear_family_g are the relations listed on construct_relation_group.
End of explanation
"""


spouse = context.one('wife', 'nuclear_family_g') | context.one('husband', 'nuclear_family_g')
printd(spouse.eval())

"""
Explanation: We can now define sets of relations the same way that we can define sets of entities, for example
End of explanation
"""


printd(henry8.follow(spouse).eval())
printd(queen_vic.follow(spouse).eval())

"""
Explanation: NQL lets you follow a group of relations, much as you can follow a single relation.
End of explanation
"""


printd(queen_vic.follow(context.all('nuclear_family_g')).eval())

"""
Explanation: As another example, these are all of the members of Queen Victoria's nuclear family.
End of explanation
"""


trip = queen_vic.follow('nuclear_family_g_subj',-1) & spouse.follow('nuclear_family_g_rel',-1)
printd(trip.eval())
print(context.get_id(list(trip.eval().keys())[0], 'nuclear_family_g_triple_t'))

t = context.one('nuclear_family_g_trip_14295', 'nuclear_family_g_triple_t')
print('triple t:')
printd(t.eval())
print('subject of t:')
printd(t.nuclear_family_g_subj().eval())
print('object of t:')
printd(t.nuclear_family_g_obj().eval())
print('relation of t:')
printd(t.nuclear_family_g_rel().eval())


"""
Explanation: Under the hood, this is implemented using the other new type, the triple type.  Let's look at a particular triple, which one which encodes the fact Queen Victoria's husband is Prince Albert.
End of explanation
"""


queen_vics_spouse_triples = queen_vic.nuclear_family_g_subj(-1) & spouse.nuclear_family_g_rel(-1)
queen_vics_spouse = queen_vics_spouse_triples.nuclear_family_g_obj()

print(queen_vics_spouse.eval())

"""
Explanation: This lets you follow relation sets using NQLs other primitives - for instance, to follow the spouse relation for Queen Victoria, we first find triples that have her as the subject, and something in the set spouse relations as the relation.  We can  then step from these triples (in this case, there's only) to the object.
End of explanation
"""


henry8_spouse_triples = henry8.nuclear_family_g_subj(-1) & spouse.nuclear_family_g_rel(-1)
print('triples:')
printd(henry8_spouse_triples.eval())
print('wives:')
printd(henry8_spouse_triples.nuclear_family_g_obj().eval())

"""
Explanation: Similarly, for Henry VIII:
End of explanation
"""


def son_fun(x): return x.son()

printd(son_fun(queen_vic).eval())

"""
Explanation: Subclassing NQLExpression
Sometimes it's clearer to define methods that you can chain, rather than functions that you call.  For instance above we added 
def child(x): return x.son() | x.daughter()
def parent(x): return x.mother() | x.father()
def grandfather(x): return parent(x).father()
allowing you to write
print child(queen_vic).eval()
print grandfather(henry8).eval()
But it might be nicer to be able use the same syntax for both user-defined things like child and KG-defined things like son.  If course, we could easily define a functional version of the KG-based relations,  like so:
End of explanation
"""


class MyNQLExpression(nql.NeuralQueryExpression):
  def child(self): return self.son() | self.daughter()

#load a new copy of the context object
context2 = nql.NeuralQueryContext()
context2.expression_factory_class = MyNQLExpression
relations = [
  'aunt', 'brother', 'daughter', 'father', 'husband', 'mother', 
  'nephew', 'niece', 'sister', 'son', 'uncle', 'wife' 
]
for r in relations:
  context2.declare_relation(r, 'person_t', 'person_t')
with open('/content/language_repo/language/nql/demos/data/royal92/royal_family.tsv') as f:
  context2.load_kg(files=f)

# now let's try using my new chaining-style method
queen_vic2 = context2.one('Victoria of house of Hanover', 'person_t')
printd(queen_vic2.child().eval())


"""
Explanation: Defining chaining-style operations is a little more complicated.  We need to subclass the NQLExpression class and add some new methods.  Then, because the context is a factory class, which generates expressions, we need to tell it to create instances of my subclass.  To avoid confustion I'll do this with a fresh context object.
End of explanation
"""


e = context.all('person_t').eval()
printd(e, 10)


"""
Explanation: Under the hood - NQL and TensorFlow
Weighted sets and counting paths
To understand the non-uniform weights that occur when we look for all the male or female people, we need to drill down into what's actually happening here as we evalulate code like. 
x = context.all('person_t')
x.mother().son() &amp; x
Let's take it step by step.  The all sets are uniformly weighted...
End of explanation
"""


e = context.all('person_t').mother().eval()
printd(e, 10)

"""
Explanation: ...but when we take the next step, we end up with differing weights:
End of explanation
"""


mothers = context.all('person_t').mother().eval()
for name,weight in sorted(mothers.items())[0:10]:
    print('%s has %d kid(s):' % (name, weight))
    printd(child(context.one(name,'person_t')).eval())

"""
Explanation: The weights here correspond to the number of children each mother has, which we can check as follows (for a sample of the values)
End of explanation
"""


mothers = context.all('person_t').mother()
print('moms:')
printd(mothers.eval(), 10)

grandmothers = mothers.mother()
print('grandmoms:')
printd(grandmothers.eval(), 10)

"""
Explanation: What's happening is that while the set of keys is the set of all mothers, the weight for any person x indicate the number of ways that x can be "demonstrated to be a mother" --- i.e., when we started with the set S and followed the relation mother, the weights for x became the number of elements in S from which x could be reached via the mother relation.
This behavior cascades - for example, if we evaluated S.mother().mother() the weight for each x' would be the number of people for whom x' is a maternal grandmother, and we will get more extreme weights.
End of explanation
"""


m = mothers.eval(as_top=6)
print('moms:')
printd(m)

"""
Explanation: In this colab, we implemented a special printd method to show the top weights in an evaluated NQL expression, so you would know exactly what was happening.  However, getting the top-weighted elements of an NQL entity set is common enough that you can do it with a special option to the eval method.
End of explanation
"""


weighted_men = context.all('person_t').mother().son()
weighted_men_tensor = weighted_men.tf
men_tensor = tf.clip_by_value(weighted_men_tensor, 0.0, 1.0)
men = context.as_nql(men_tensor,'person_t')

printd(men.eval(), 10)

"""
Explanation: While you can ignore the weights ,  the NQL implementors believes that this behavior is actually a feature and not a bug, as the weights often are useful information.  Here, for example, the weights reflect the number of descendants.
If you want to get rid of the  non-uniform weights, one way would be to convert the weighted NQL expression to Tensorflow - which can be done with the .tf property -- clip the weights using the appropriate Tensorflow op, and then convert back to NQL,  as follows.
End of explanation
"""


men2 = weighted_men.tf_op(lambda t:tf.clip_by_value(t, 0.0, 1.0))
printd(men2.eval(), 10)

"""
Explanation: If you're just jumping out to Tensorflow for a short time, it's sometimes more concise to use the tf_op method, which lets you use any Python function that transforms a tensor to one with the same shape.
End of explanation
"""


m = henry8.wife().eval(as_dicts=False)
print('m', m)
print('m.shape', m.shape)
print('nonzeros in m', np.nonzero(m))


"""
Explanation: We discuss the possible Tensorflow - NQL conversion processes below.
Following relations via sparse matrix multiplication
Under the hood, what's going on here is that each weighted set is represented internally as a vector - generally a sparse vector, with one component for each entity.   To make it a little more complicated, NQL is, like most neural systems, set up to operate on minibatches, so the internal representation of a single set is a matrix.  For these queries, the matrix has one row, and n columns, where n is the number of person_t entiities.  A relation is represented as a matrix, and following a relation is implemented as a vector-matrix multiplication.
Let's drill down into this a little more.
There are two equivalent ways to get at this 'raw' representation.  One is to pass the as_dicts=False option into the eval method for an NQL expression.
End of explanation
"""


m = henry8.wife().tf.numpy()
print('m', m)
print('m.shape', m.shape)
print('nonzeros in m', np.nonzero(m))

"""
Explanation: Alternatively we can get at the underlying tensorflow expression, and evaluate it:
End of explanation
"""


row_indices, col_indices = np.nonzero(henry8.wife().tf.numpy())
for person_id in col_indices.tolist():
  print('person with id %d is %s' % (person_id, context.get_entity_name(person_id, 'person_t')))
# double check these values are correct
print('target set is:')
printd(henry8.wife().eval().keys())

"""
Explanation: In this representation, the nonzero column indices are ids for the wives of Henry VIII.  We can decode these ids using the context object.  (The row indices are all zero, since the matrix has only one row.)
End of explanation
"""


printd(context.as_nql(m, 'person_t').eval())

"""
Explanation: As an alternative, we could convert m back to an NQL expression to examine it.
End of explanation
"""


wife_tensor = context.get_tf_tensor('wife')
print(wife_tensor)

henry8_id = context.get_id('Henry_VIII of house of Tudor', 'person_t')
for row_index,col_index in wife_tensor.indices.numpy().tolist():
  if col_index==henry8_id:
      print('id: %d name: %s' % (row_index, context.get_entity_name(row_index, 'person_t')))

"""
Explanation: Every relation is represented as a sparse matrix.  Below we'll look at the sparse matrix for the wife relation, and decode a few indices as a check.
End of explanation
"""


v = henry8.eval(as_dicts=False)
S = wife_tensor
henrys_wives = tf.transpose(tf.sparse.sparse_dense_matmul(S, tf.transpose(v)))
print(np.nonzero(henrys_wives.numpy()))

"""
Explanation: Aside: The storage is a little counterintuitive: instead of wife_matrix mapping from Henry to Ann Bolyn, for instance, by making the husbands the row indices and the wives the column indices, NQL stores indices in the opposite order.  This is because Tensorflow doesn't support multiplying a dense vector v  by a sparse matrix S, only multiplication of S by v.  So instead of computing v * S, we must compute something more like
transpose( transpose(S) * transpose(v) )
Hence storing transpose(S) directly is somewhat more efficient.
So finally, if we want to simulate henry8.wife(), we can do it as follows.
End of explanation
"""


