from pyspark import SparkConf
from pyspark import SparkContext
from pyspark.mllib.feature import Word2Vec

import re

from urllib.request import urlopen

from functools import reduce

from pickle import dumps as pds

import pymongo

def cleanstr(s):
    noPunct = re.sub("[^a-z ]", " ", s.lower())
    collapsedWhitespace = re.sub("(^ )|( $)", "", re.sub("  *", " ", noPunct))
    return collapsedWhitespace


def url2rdd(sc, url):
    response = urlopen(url)
    corpus_bytes = response.read()
    text = str(corpus_bytes).replace("\\r", "\r").replace("\\n", "\n")
    rdd = sc.parallelize(text.split("\r\n\r\n"))
    rdd.map(lambda l: l.replace("\r\n", " ").split(" "))
    return rdd.map(lambda l: cleanstr(l).split(" "))


def train(sc, urls):
    w2v = Word2Vec()
    rdds = reduce(lambda a, b: a.union(b), [url2rdd(sc, url) for url in urls])
    return w2v.fit(rdds)


def workloop(master, inq, outq, dburl):
    sconf = SparkConf().setAppName("ophicleide-worker").setMaster(master)
    sc = SparkContext(conf=sconf)

    if dburl is not None:
        db = pymongo.MongoClient(dburl).ophicleide

    outq.put("ready")

    while True:
        job = inq.get()
        urls = job["urls"]
        mid = job["_id"]
        model = train(sc, urls)

        words, vecs = zip(*[(w, list(v)) for w, v in model.getVectors().items()])

        # XXX: do something with callback here
        
        if dburl is not None:
            db.models.update_one({"_id": mid}, {"$set": {"status": "ready", "model": {"words": list(words), "vecs": list(vecs)}}, "$currentDate": {"last_updated": True}})

        outq.put((mid, name))

@@ -181,55 +181,63 @@
 ce%0A%0A
-from pickle import dumps as pds%0A%0Aimport pymongo
+import pymongo%0A%0Aimport zlib%0A%0Afrom numpy import ndarray%0A
 %0A%0Ade
@@ -1381,24 +1381,16 @@
 ck here%0A
-        
 %0A       
@@ -1404,32 +1404,323 @@
 rl is not None:%0A
+            ndvecs = ndarray(%5Blen(words), len(vecs%5B0%5D)%5D)%0A            for i in range(len(vecs)):%0A                ndvecs%5Bi%5D = vecs%5Bi%5D%0A%0A            ns = ndvecs.dumps()%0A            zns = zlib.compress(ns, 9)%0A%0A            print(%22len(ns) == %25d; len(zns) == %25d%22 %25 (len(ns), len(zns)))%0A            %0A
             db.m
@@ -1816,25 +1816,21 @@
 ), %22
+znd
 vecs%22: 
-list(vecs)
+zns
 %7D%7D, 
@@ -1898,11 +1898,18 @@
 id, 
+job%5B%22
 name
+%22%5D
 ))%0A

store word vectors as compressed ndarray