##
##  Nice demo of nearest-neighbour search (on verb-object cooccurrence matrix from ukWaC corpus)
##

##
## 1) Setup (read prepared SVD model, 300 dimensions)
##

library(MASS) # needed for visualisation below

# DSM based on syntactic verb-object cooccurrences in ukWaC corpus (ca. 2 billion words).
#  - targets = nouns (f >= 20),  features = verbs (f >= 500),  both lemmatised
#  - feature weights = log-scaled t-score,  row vectors normalised to Euclidean unit length
#  - SVD dimensionality reduction to 300 dimensions,  row vectors renormalised

load("ukwac_vobj_S_svd.rda") # R data format is compact and can be loaded efficiently
S <- ukwac.vobj$svd300$M     # extract DSM matrix


##
## 2) Utility functions: cosine distance & nearest neighbours (with visualisation)
##

# calculate p-norms of rows in case cosine() is used on an unnormalised matrix
p.norm <- function (M, p=2) {
	stopifnot(p >= 1)
	if (p == 1) {
		.norm <- function (x) sum(abs(x))
	} else {
		.norm <- function (x) (sum(abs(x)^p))^(1/p)
	}
	return(apply(M, 1, .norm))
}

# compute cosine similarity (or angles) between rows of matrices
cosine <- function (M, M2=M, angles=FALSE, normalised=FALSE) {
	sim <- M %*% t(M2) # should calculate dot products between rows of M and rows of M2
	if (!normalised) {
		norms.M <- p.norm(M, 2) # norms of row vectors (if not normalised yet)
		norms.M2 <- p.norm(M2, 2)
		sim <- sim / outer(norms.M, norms.M2)
	}
	if (angles) {
		stopifnot(all(sim >= -(1+1e-6) & sim <= 1+1e-6)) # check that cosines are in appropriate range
		sim <- pmin(pmax(sim, -1), 1) # clamp to range [-1, 1]
		sim <- acos(sim) / pi * 180 # angles are returned in degrees
	}
	rownames(sim) <- rownames(M)
	colnames(sim) <- rownames(M2)
	return(sim)
}


##
## 3) Utility functions: nearest neighbours with optional visualisation
##

neighbours <- function (term, n=10, M=S, plot=c("none", "pca", "mds", "sammon"), n.plot=n, plot.range=NULL, iterate=FALSE) {
	plot <- match.arg(plot)
	if (! term %in% rownames(M)) stop("Noun '", term, "' is not in database.")
	term.vec <- S[term,,drop=FALSE] # make sure vector is interpreted as single-row matrix
	d <- cosine(M, term.vec, normalised=TRUE, angles=TRUE)
	idx <- order(d)
	idx.show <- idx[2:(n+1)]  # first "neighbour" is term itself at distance 0
	if (plot != "none") {
		if (n.plot < 5) iterate <- FALSE
		iterate.n <- if (iterate) 5:n.plot else n.plot
		coords <- NULL
		for (n.plot in iterate.n) {
			idx.plot <- idx[1:(n.plot+1)] # map term vector and n.plot nearest neighbours into 2D plane
			N <- M[idx.plot, ]
				if (plot == "pca") {
				P <- prcomp(N, center=TRUE, scale=FALSE)
				coords <- P$x[, 1:2]
			} else {
				N.d <- cosine(N, normalised=TRUE, angles=TRUE)
				if (iterate && !is.null(coords)) {
					# add dummy entry for new data point, then shake well
					coords <- jitter(rbind(coords, c(0,0)), amount=10) 
					coords <- sammon(N.d, y=coords, k=2, trace=FALSE)$points
				} else {
					if (plot == "sammon" || iterate) {
						coords <- sammon(N.d, k=2, trace=FALSE)$points
					}
					else {
						coords <- isoMDS(N.d, k=2, p=2, trace=FALSE)$points
					}
				}
			}
			labels <- rownames(M)[idx.plot]
			if (is.null(plot.range)) {
				x.range <- extendrange(coords[, 1], f=.05)
				y.range <- extendrange(coords[, 2], f=.05)
			} else {
				if (length(plot.range) == 1) plot.range <- c(-plot.range, plot.range)
				x.range <- y.range <- plot.range
			}
			if (iterate) Sys.sleep(.2)
			plot(coords, pch=20, xlim=x.range, ylim=y.range, xlab="", ylab="")
			points(coords[1,1], coords[1,2], pch=20, cex=1.2, col="red")
			text(coords[-1, ], labels=labels[-1], pos=3, font=2)
			text(coords[1,1], coords[1,2], labels=labels[1], pos=3, font=2, cex=1.2, col="red")
		}
	}
	round(d[idx.show, ], 2) 
}

NN <- function (term, n=15, plot="mds", n.plot=n, range=NULL, iterate=FALSE) {
	dT <- system.time(term.NN <- neighbours(term, n=n, plot=plot, n.plot=n.plot, plot.range=range, iterate=iterate))
	cat("[",round(dT[3], 3)," seconds]\n", sep="")
	term.NN <- t(t(term.NN))
	colnames(term.NN) <- "degrees"
	print(term.NN)
}
NNi <- function (term, n=30, n.disp=30, range=NULL, plot="sammon") {
	NN(term, n=n.disp, plot=plot, n.plot=n, range=range, iterate=TRUE)
}


##
## 4) Usage examples
##

neighbours("beer") # returns vector with 10 nearest neighbours and their angular distance
neighbours("beer", 30) # 30 neighbours

par(cex=1.2, mar=c(2,2,.5,.5)) # set up plot window for visualisation

neighbours("beer", 20, plot="pca") # PCA visualisation of 20 nearest neighbours
neighbours("beer", 50, plot="pca") # large neighbourhoods get a bit crowded
neighbours("beer", 50, plot="mds") # MDS visualisation (multi-dimensional scaling)
neighbours("beer", 50, plot="sammon") # another MDS algorithm

# specify a fixed plotting range to produce comparable plots (depends on algorithm!)
par(mfrow=c(2,2))
neighbours("beer", 20, plot="mds", plot.range=30)
neighbours("coffee", 20, plot="mds", plot.range=30)
neighbours("boy", 20, plot="mds", plot.range=30)
neighbours("marginality", 20, plot="mds", plot.range=30)
par(mfrow=c(1,1))

# iterative plot: add one neigbour at a time as simple animation (use with "sammon" method)
neighbours("beer", 30, plot="sammon", plot.range=35, iterate=TRUE)
neighbours("beer", 30, plot="sammon", plot.range=35) # result is often different from non-iterative plot

# if you want to demonstrate this word space interactively, use NN() and NNi() functions
NN("love")  # 15 nearest neighbours with MDS plot
NN("love", plot="pca")
NNi("love") # 30 nearest neighbours with incremental animation
