# -*- coding: utf-8 -*-
"""Wrapper for a file containing brown clusters of a corpus."""
from __future__ import absolute_import, division, print_function, unicode_literals

class BrownClusters(object):
    """
    Wrapper to load and handle and file with brown clusters, as generated by the tool of Percy Liang
    on https://github.com/percyliang/brown-cluster .
    The tool generates multiple files. Only the file 'paths' is used here.
    Example usage:
        bc = BrownClusters("/some/directory/paths")
        cluster_idx = bc.get_cluster_of("foo")
    """
    def __init__(self, filepath):
        """Initialize the class.
        Args:
            filepath: The filepath to the file 'paths' file containing the brown clusters.
        """
        self.word_to_cluster = dict()
        self.word_to_bitchain = dict()
        self.fill_from_file(filepath)

    def clear(self):
        """Reset this class, deletes all word->cluster and word->bitchain mappings."""
        self.word_to_cluster = dict()
        self.word_to_bitchain = dict()

    def fill_from_file(self, filepath):
        """Loads a 'paths' file with brown clusters as generated by Percy Liang's tool.
        After loading, this class will contain all word->cluster and word->bitchain mappings of
        the file.

        Example content of such a file (looks like a cluster of names):
            0001000	Musil	701
            0001000	Hackl	701
            0001000	Focke	702
            0001000	Gundlach	702
            0001000	Hornung	702
            0001000	Harlow	702
            0001000	Bartholomew	702
            0001000	McCall	702
            0001000	Bloomfield	703

        Note: This function is automatically called in __init__().

        Args:
            filepath: The filepath to the file 'paths' file containing the brown clusters.
        """
        with open(filepath, "r") as handle:
            last_count = -1
            cluster_idx = 1

            for line_idx, line in enumerate(handle):
                columns = line.decode("utf-8").strip().split("\t")
                if len(columns) == 3:
                    bitchain, word, count = columns
                    count = int(count)

                    if count < last_count:
                        cluster_idx += 1
                    last_count = count

                    self.word_to_cluster[word] = cluster_idx
                    self.word_to_bitchain[word] = bitchain
                else:
                    print("[Warning] Expected 3 columns in brown clusters file at line %d, " \
                          "got %d" % (line_idx, len(columns)))

    def get_cluster_of(self, word, default=-1):
        """Returns the brown cluster of a word.

        Args:
            word: The word.
            default: A default value to return if the word was not contained in the file.
        Returns:
            cluster id (integer)
            or provided default value, if the word was not contained in the file.
        """
        if word in self.word_to_cluster:
            return self.word_to_cluster[word]
        else:
            return default

    def get_bitchain_of(self, word, default=""):
        """Returns the bitchain of a word.

        The bitchain would be the first column in the following example excerpt:
            0001000	Base	6857
            0001000	Rosenberg	6885
            0001000	Shakespeare	6891
            0001000	Schumacher	7007
        The bitchain resembles the position of a cluster in a tree of all brown clusters.
        Deeper clusters start with the same bitchain.

        Args:
            word: The word.
            default: A default value to return if the word was not contained in the file.
        Returns:
            cluster id (integer)
            or provided default value, if the word was not contained in the file.
        """
        if word in self.word_to_bitchain:
            return self.word_to_bitchain[word]
        else:
            return default
