#!/usr/bin/env python

"""
Cluster words from the Switchboard Dialog Act Corpus (from a supplied
set of syntactic categories) using k-means clustering.

If this program is from the command line (calling its __main__
method), then it performs the clustering for the single-tag set ['uh']
(interjection), because the words with that tag are defined perhaps
entirely by their usage conditions in dialog. Thus, this demo provides
a nice (successful) illustration of the method.

(To run the main-method, this script must be in the same directory
as swda.py and the swda corpus root directory.)

---Chris Potts
"""

######################################################################

from collections import defaultdict
from operator import itemgetter
import numpy
from nltk import cluster # the clustering API
import nltk.cluster.util # contains the distance measures
from swda import CorpusReader

######################################################################

class SwdaWordTagClusterer:
    """
    For a restricted set of Penn Treebank tags, cluster words with
    those tags according to their distribution relative to the DAMSL
    dialog act tags.  We build a matrix

          act_tag1 act_tag2 act_tag3 ... act_act43
    word1
    word2
    word3
    ...
    wordN

    (N = the size of the vocab), where the cells contain
    length-normalized relative frequency values. The idea is that the
    rows represent word-meanings as defined by co-occurrence with the
    act tags. We apply k-means clustering to these vectors using the
    NLTK clustering functionality.
    """
    def __init__(self, cats, corpus, count_threshold=20,
                 num_means=2, repeats=1,
                 distance_measure=nltk.cluster.util.euclidean_distance):
        """
        Arguments:
        cats -- a list of Penn Treebank word-level tags (case ignored)
        corpus -- an swda.CorpusReader object
        count_threshold -- exclude words with token-counts below this level (default: 20)
        num_means -- the number of clusters to return (default: 2)
        repeats -- the number of times to repeat clustering, with a new random seed each time
        distance_measure -- the distance measure used by the clusterer (default: nltk.cluster.util.euclidean_distance)
        """        
        self.cats = map(str.lower, cats) # downcase to match the wn-lemmatizing style of swda.CorpusReader
        self.corpus = corpus
        self.count_threshold = count_threshold
        self.num_means = num_means
        self.repeats = repeats
        self.distance_measure = distance_measure
        # Create the count dictionary word --> DAMSL-tag --> count:
        self.count_dict = self.build_count_dictionary()        
        # Create the count matrix:
        self.mat, self.vocab, self.all_tags = self.build_matrix()
        # Normalize self.mat by length:
        self.length_normalize_matrix()
       
    def build_count_dictionary(self):
        """
        Build the count distribution word --> DAMSL-tag --> count
        restricting attention to word tokens tagged with one of the
        categories in self.cats.
        Value:
        d  -- a two-dimensional dict with default value 0
        """
        i = 0
        d = defaultdict(lambda : defaultdict(int))        
        for utt in corpus.iter_utterances():
            for word, pos in utt.pos_lemmas(wn_lemmatize=True):
                if pos in cats:
                    for tag in utt.damsl_act_tags():
                        d[word.lower()][tag] += 1
            # This institutes a short run for basic error checking:
            i += 1
            if i == 5000: return d
        # Impose the count threshold.
        for word, tag_dict in d.items():
            if sum(tag_dict.values()) < self.count_threshold:
                del d[word]
        return d

    def build_matrix(self):
        """
        Create a two-dimensional length-normalized relative frequency matrix.    
        Value:
        mat -- a two-dimensional (n, m) numpy array of relative
               frequencies, where n is the vocab size and m is the
               number of tags               
        vocab -- a sorted list of vocab items (strings), where the
                 indices correspond to those of the rows in mat
        """
        # Sorted list of all the tags:
        all_tags = sorted(list(set([tag for tag_dist in self.count_dict.values() for tag in tag_dist.keys()])))
        # Sorted list of all the vocab items:
        vocab = sorted(self.count_dict.keys())
        # Initialize an all-zeros matrix with the right dimensions: words
        # as rows, tags as columns.
        mat = numpy.zeros((len(vocab), len(all_tags)))
        # Fill the matrix with relative frequencies 
        for i in xrange(len(vocab)):
            for j in xrange(len(all_tags)):
                mat[i,j] = self.count_dict[vocab[i]][all_tags[j]]
        return (mat, vocab, all_tags)

    def length_normalize_matrix(self):
        """Length-normalize the row vectors of self.mat."""
        def vec_norm(vec):
            return vec / numpy.sqrt(numpy.dot(vec, vec))
        self.mat = map(vec_norm, self.mat)

    def kmeans(self):
        """
        Use the the NLTK kmeans clustering functions to cluster the
        word-vectors (which are length-normalized by the clusterer).        
        Output:
        clusters (defaultdict(list)) -- a mapping from cluster indices to lists of strings
        """
        # Set-up the clusterer using the user's parameters (num_means = num clusters),
        # and specify that that the vectors should be length-normalized so that we
        # avoid clustering by overall frequency:
        clusterer = cluster.KMeansClusterer(self.num_means, self.distance_measure, repeats=self.repeats, normalise=False)
        cluster_vector = clusterer.cluster(self.mat, assign_clusters=True, trace=False)
        # Build a mapping from cluster indices to the words in them:
        clusters = defaultdict(list)
        for i in xrange(len(self.vocab)):
            word = self.vocab[i]
            cluster_index = cluster_vector[i]
            clusters[cluster_index].append(word)        
        return clusters

######################################################################

def length_normalization(vec):
    return vec / numpy.sqrt(numpy.dot(vec, vec))

def matrix_inspection(clust):
    import csv
    mat = map(length_normalization, clust.mat)
    csvwriter = csv.writer(file('tag-matrix-uh-normed.csv', 'w'))
    csvwriter.writerow(['Word'] + clust.all_tags)
    for i in xrange(len(mat)):
        csvwriter.writerow([clust.vocab[i]] + list(mat[i]))
        
if __name__ == '__main__':
    """ Run an example clustering experiment for cats."""
        
    cats = ['uh']
    # cats = ['prp', 'prp$', 'wp', 'wp$']
    # cats = ['in']
    corpus = CorpusReader('swda')
    clust = SwdaWordTagClusterer(cats, corpus, count_threshold=20, num_means=5, repeats=10, distance_measure=nltk.cluster.util.euclidean_distance)
    clusters = clust.kmeans()
    print "======================================================================"
    for cluster_index, words in clusters.items():
        print cluster_index, words
    print "======================================================================"