#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
"""Test."""

from sklearn.cluster import KMeans
from gensim.models import Word2Vec
import os
import cPickle
import codecs
import logging
import json
import numpy as np


def cluster(language, vec_list, word_list, num_clusters=1000, n_jobs=3):
    """Utilised for clustering passed Vector List, and return n clusters."""
    kmeans_clust = KMeans(
        n_clusters=num_clusters,
        precompute_distances=True,
        n_jobs=n_jobs,
    )
    idx = kmeans_clust.fit_predict(vec_list)
    k = [[] for _ in xrange(num_clusters)]
    word_centroid_map = dict(zip(word_list, idx))
    for word in word_centroid_map.keys():
        k[word_centroid_map[word]].append(word)
    with codecs.open(os.path.join(
        'data', 'vectors', '%s-%s-clustering.json' % (language, num_clusters)
    ), 'w', encoding='utf-8') as file:
        file.write(
            json.dumps(k)
        )

if __name__ == '__main__':
    language = 'english'
    num_clusters = 1000
    logging.basicConfig(level=logging.DEBUG)
    model = Word2Vec.load(
        os.path.join('data', 'models', 't-vex-english-model')
    )
    word_list = np.array(model.vocab.keys())
    vec_list = np.array([model[word] for word in word_list])
    del(model)
    with codecs.open(
        os.path.join(
            'data', 'vectors', '%s-%s-vectors' % (language, num_clusters)
        ), 'w'
    ) as file:
        file.write(
            cPickle.dumps({
                'vector_list': vec_list,
                'word_list': word_list,
            })
        )
    cluster(
        language=language,
        vec_list=vec_list,
        word_list=word_list,
        num_clusters=num_clusters
    )

@@ -49,13 +49,51 @@
 %0A%22%22%22
-Test.
+**Module generates clusters from a Model.**
 %22%22%22%0A
@@ -330,16 +330,23 @@
 %0A    %22%22%22
+%0A    **
 Utilised
@@ -403,16 +403,487 @@
 lusters.
+**%0A%0A    **API Documentation:**%0A        :param language: Language utilised for filename for cluster file.%0A        :param vec_list: A vector list obtained from the Model.%0A        :param word_list: A word list obtained from the Model corresponding to vec_list.%0A        :param num_clusters: Num of clusters to be generated %5B Default 1000 %5D%0A        :param n_jobs: Num of jobs %5B multiprocessing %5D%0A%0A    .. seealso::%0A        * :mod:%60modules.model_generator.model_generation%60%0A    
 %22%22%22%0A    

Update Documentation to rst for automatic autodoc support of Sphinx.