# -*- coding: utf-8 -*-
"""Encapsulates handling of a word2vec clusters file."""
from __future__ import absolute_import, division, print_function, unicode_literals

class W2VClusters(object):
    """Encapsulates handling of a word2vec clusters file.
    The file can be generated with the word2vec tool using the flag "-classes"."""
    def __init__(self, filepath):
        """Initializes a new W2VClusters object.
        Args:
            filepath: Filepath to the file containing the w2v clusters.
        """
        self.word_to_cluster = dict()
        self.fill_from_file(filepath)

    def clear(self):
        """Resets this object, i.e. empties the dictionary."""
        self.word_to_cluster = dict()

    def fill_from_file(self, filepath):
        """Fills the object's dictionary (mapping word to cluster) from a file.
        The file must be generated by the word2vec tool (add argument -classes) or follow
        the same structure, which is
            foo 123
            bar 123
            bla 76
            asd 76
        i.e. first column is the word, second column is the cluster, both are separated by a
        whitespace.

        Args:
            filepath: Filepath to the word2vec clusters file.
        """
        with open(filepath, "r") as handle:
            for line_idx, line in enumerate(handle):
                columns = line.decode("utf-8").strip().split(" ")
                if len(columns) == 2:
                    word = columns[0]
                    cluster_idx = int(columns[1])
                    self.word_to_cluster[word] = cluster_idx
                else:
                    print("[Warning] Expected 2 columns in w2v clusters file at line %d, " \
                          "got %d" % (line_idx, len(columns)))

    def get_cluster_of(self, word, default=-1):
        """Returns the w2v cluster of a word.
        Args:
            word: The word for which to return the cluster.
            default: A default value to return if the word was not found in the w2v file.
        Returns:
            integer or default value (-1).
        """
        if word in self.word_to_cluster:
            return self.word_to_cluster[word]
        else:
            return default
