# Jorge Castanon, February 2016
# Data Scientist @ IBM

# run in terminal with:
# ~/Documents/spark-1.5.1/bin/spark-submit mllib-scripts/visualize-words.py
# Replace this line with:
# /YOUR-SPARK-HOME/bin/spark-submit mllib-scripts/visualize-words.py

import numpy as np
import math

from pyspark.context import SparkContext

# next 2 lines can be replaced to read from hdfs, 
# if the Word2Vec matrix is big 
Feat = np.load('mllib-scripts/myW2Vmatrix.npy')    # reads model generated by Word2Vec
words = np.load('mllib-scripts/myWordList.npy')    # reads list of words
wordToModel = 'christmas'
maxWordsVis = 30

Feat = np.load('models/w2v_may1_may19_june1_june11.npy')
words = np.load('models/word_may1_may19_june1_june11.npy')
wordToModel = 'data'
maxWordsVis = 10

print "\n================================================="
print "Size of the Word2Vec matrix is: ", Feat.shape 
print "Number of words in the models: ", words.shape
print "=================================================\n"

## Spark Context
sc = SparkContext('local','visualize-words') 

## Read the Word2Vec model
# the next line should be read/stored from hdfs if it is large
Feat = sc.parallelize(Feat) 

# map feature matrix to spark vectors
from pyspark.mllib.linalg import Vectors
Feat = Feat.map(lambda vec: (Vectors.dense(vec),))

## Define a df with feature matrix
from pyspark.sql import SQLContext
sqlContext = SQLContext(sc)
dfFeat = sqlContext.createDataFrame(Feat,["features"])
dfFeat.printSchema()

## PCA to project Feature matrix to 2 dimensions
from pyspark.ml.feature import PCA
numComponents = 3
pca = PCA(k=numComponents, inputCol="features", outputCol="pcaFeatures")
model = pca.fit(dfFeat)
dfComp = model.transform(dfFeat).select("pcaFeatures")
# get the first two components to lists to be plotted

compX = dfComp.map(lambda vec: vec[0][0]).take(maxWordsVis)
compY = dfComp.map(lambda vec: vec[0][1]).take(maxWordsVis)
compZ = dfComp.map(lambda vec: vec[0][2]).take(maxWordsVis)

## finish Spark session
sc.stop()

## plot
fs=20 #fontsize
w = words[0:maxWordsVis]
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(compX, compY, compZ, color='red', s=100, marker='o', edgecolors='black')
for i, txt in enumerate(w):
	ax.text(compX[i],compY[i],compZ[i], '%s' % (txt), size=fs-5, zorder=1, color='k')
ax.set_xlabel('1st. Component', fontsize=fs)
ax.set_ylabel('2nd. Component', fontsize=fs)
ax.set_zlabel('3rd. Component', fontsize=fs)
title = 'Top',str(maxWordsVis),'closest words to term "',wordToModel ,'" in Twitter'
ax.set_title(' '.join(title), fontsize=fs)
ax.grid(True)
plt.show()

