%matplotlib inline
from __future__ import print_function
import os
from pyspark import SQLContext
from pyspark.sql import Row
import pyspark.sql.functions as sql
import pyspark.sql.types as types
#from pyspark.sql.functions import udf, length
import matplotlib.pyplot as plt
import numpy
import math
import matplotlib.pyplot as plt
import seaborn as sns

import nltk
import pyspark.ml.feature as feature

# Load Processed Parquet
sqlContext = SQLContext(sc)
notes = sqlContext.read.parquet("../data/idigbio_notes.parquet")
total_records = notes.count()
print(total_records)
# Small sample of the df
notes = notes.sample(withReplacement=False, fraction=0.1)
notes.cache()
print(notes.count())

# Still have some problems with the document field having nulls and
# being empty. Not sure where nulls came from but likely the 
# empties are really whitespace
#notes = notes.select(sql.trim(notes["document"]).alias("document"))\
#    .dropna(subset="document")\
#    .filter(sql.length("document") > 0)


#print(notes.count())

# Moved the above to the text_characterization notebook and re-wrote parquet
# should be clean now.

notes.select(notes["document"])\
    .orderBy(notes["document"])\
    .show(10, truncate=False)
notes.select(notes["document"])\
    .orderBy(notes["document"], ascending=False)\
    .show(10, truncate=False)

"""
Explanation: Identifying relationships in notes
Goal of this notebook is to determine how successful we can find relationships
End of explanation
"""


from lib.tokens import Tokens
tokens = Tokens()
t = tokens.tokenize("Hello, my name is Mace Windoo")
print(t)

udf_tokenize = sql.udf(tokens.tokenize, types.ArrayType(types.StringType()))

notes_w_tokens = notes.withColumn('tokens', udf_tokenize(notes['document']))
for r in notes_w_tokens.select(notes_w_tokens["tokens"]).head(50):
    print(" | ".join(r["tokens"]))
    print("\n")

notes_w_tokens.printSchema()

from lib.pos_tags import PosTags
pos_tags = PosTags()

pos = pos_tags.tag(t)
print(pos)
print(type(pos))
print(type(pos[0]))
print(type(pos[0]["tag"]))
print(pos[0])

udf_part_of_speech = sql.udf(pos_tags.tag, types.ArrayType(
                                    types.MapType(
                                        types.StringType(),
                                        types.StringType()
                                    )
                                )
                            )

notes_w_tokens2 = notes_w_tokens.withColumn('pos', 
                                            udf_part_of_speech(notes_w_tokens['tokens']))

for r in notes_w_tokens2.select(notes_w_tokens2["pos"]).head(50):
    s = ""
    for p in r["pos"]:
        s = s + p["word"] + " (" + p["tag"] + ") | "
    print(s + "\n")

notes_w_tokens2.printSchema()

# Can we work with maps natively?
notes_w_tokens2.select(notes_w_tokens2["pos"][0]["word"]).show(3, truncate=False)
# YES!

# Let's go pipline our whole analysis
from lib.tokens import Tokens
from lib.pos_tags import PosTags
from lib.relations import Relations

t = Tokens()
p = PosTags()
r = Relations()

def pipeline(s):
    '''
    Given a string, return a list of relations
    '''
    return r.find(p.tag(t.tokenize(s)))

pipeline_udf = sql.udf(pipeline, types.ArrayType(
                                       types.MapType(
                                               types.StringType(), 
                                               types.MapType(
                                                       types.StringType(),
                                                       types.StringType()
                                                       )
                                               )
                                       )
                    )

relations = notes.withColumn("rels", pipeline_udf(notes["document"]))
relations.cache()
relations.select(relations["rels"]).show(3, truncate=False)

# Maybe graph distribution of relations per record?


# Let's at least count up how many relations we have


"""
Explanation: Sentence detection
Does splitting in to sentences matter? Is there enough information to do this with a natural language library or should things like "," "[]", and "{}" be worked in to address semi-structured data?
Tokenize documents
End of explanation
"""


# Split out words by type
# Can't figure out how to access elements of a map in a filter so 
# build something that filters the lists for us.
def find_pos(pos, part):
    '''
    Take a list of dicts that represent words tagged with
    pos information and return a list of words that match
    the requested pos
    '''
    retval = []
    for p in pos:
        if p["tag"].startswith(part):
            retval.append(p["word"])
    return retval

print(find_pos(pos, "NN"))

# Can't figure out how to pass a single string to a UDF
find_nouns_udf = sql.udf(lambda x: find_pos(x, "NN"), types.ArrayType(types.StringType()))


nouns = notes_w_tokens2\
    .select(sql.explode(find_nouns_udf(notes_w_tokens2["pos"])).alias("word"))
nouns.cache()
nouns.show(3)

noun_counts = nouns\
    .groupBy("word")\
    .count()\
    .orderBy("count", ascending=False)\
    
noun_counts.show(30)

noun_counts.cache()

noun_counts.orderBy(noun_counts["count"]).show(30)

noun_counts_pdf = noun_counts.limit(1000).toPandas()
print(noun_counts_pdf.head())

tuples = []
for l in noun_counts_pdf.iterrows():
    tuples.append( (l[1], l[0]) )

%matplotlib inline
from wordcloud import WordCloud

wordcloud = WordCloud().generate_from_frequencies(tuples)

# And some verbs

"""
Explanation: Leave this for later
End of explanation
"""


