# Copyright 2019-2020 by Christopher C. Little.
# This file is part of Abydos.
#
# Abydos is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Abydos is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Abydos. If not, see <http://www.gnu.org/licenses/>.

"""abydos.corpus._unigram_corpus.

Unigram Corpus
"""

import pickle  # noqa: S403
from codecs import open as c_open
from collections import Counter, defaultdict
from math import log1p
from typing import Any, Callable, DefaultDict, Optional, Tuple

from ..tokenizer import _Tokenizer

__all__ = ['UnigramCorpus']


def _dd_default(*args: Any) -> Tuple[int, int]:
    return 0, 0


class UnigramCorpus:
    """Unigram corpus class.

    Largely intended for calculating inverse document frequence (IDF) from a
    large corpus of unigram (or smaller) tokens, this class encapsulates a
    dict object. Each key is a unigram token whose value is a tuple consisting
    of the number of times a term appeared and the number of distinct documents
    in which it appeared.

    .. versionadded:: 0.4.0
    """

    def __init__(
        self,
        corpus_text: str = '',
        documents: int = 0,
        word_transform: Optional[Callable[[str], str]] = None,
        word_tokenizer: Optional[_Tokenizer] = None,
    ) -> None:
        r"""Initialize UnigramCorpus.

        Parameters
        ----------
        corpus_text : str
            The corpus text as a single string
        documents : int
            The number of documents in the corpus. If equal to 0 (the default)
            then the maximum from the internal dictionary's distinct
            documents count.
        word_transform : function
            A function to apply to each term before term tokenization and
            addition to the corpus. One might use this, for example, to apply
            Soundex encoding to each term.
        word_tokenizer : _Tokenizer
            A tokenizer to apply to each sentence in order to retrieve the
            individual "word" tokens. If set to none, str.split() will be used.

        Example
        -------
        >>> tqbf = 'The quick brown fox jumped over the lazy dog.\n'
        >>> tqbf += 'And then it slept.\n And the dog ran off.'
        >>> corp = UnigramCorpus(tqbf)


        .. versionadded:: 0.4.0

        """
        self.corpus = defaultdict(
            _dd_default
        )  # type: DefaultDict[str, Tuple[int, int]]
        self.transform = word_transform
        self.tokenizer = word_tokenizer
        self.doc_count = documents

        self.add_document(corpus_text)

    def add_document(self, doc: str) -> None:
        """Add a new document to the corpus.

        Parameters
        ----------
        doc : str
            A string, representing the document to be added.


        .. versionadded:: 0.4.0

        """
        for word, count in Counter(doc.split()).items():
            self._add_word(word, count, 1)
        self.doc_count += 1

    def save_corpus(self, filename: str) -> None:
        """Save the corpus to a file.

        This employs pickle to save the corpus (a defaultdict). Other
        parameters of the corpus, such as its word_tokenizer, will not be
        affected and should be set during initialization.

        Parameters
        ----------
        filename : str
            The filename to save the corpus to.


        .. versionadded:: 0.4.0

        """
        with open(filename, mode='wb') as pkl:
            pickle.dump(self.corpus, pkl)

    def load_corpus(self, filename: str) -> None:
        """Load the corpus from a file.

        This employs pickle to load the corpus (a defaultdict). Other
        parameters of the corpus, such as its word_tokenizer, will not be
        affected and should be set during initialization.

        Parameters
        ----------
        filename : str
            The filename to load the corpus from.


        .. versionadded:: 0.4.0

        """
        with open(filename, mode='rb') as pkl:
            self.corpus = pickle.load(pkl)  # noqa: S301
        self._update_doc_count()

    def _update_doc_count(self) -> None:
        """Update document count, if necessary.

        .. versionadded:: 0.4.0
        """
        max_docs = max(self.corpus.values(), key=lambda _: _[1])[1]
        self.doc_count = max(max_docs, self.doc_count)

    def _add_word(self, word: str, count: int, doc_count: int) -> None:
        """Add a term to the corpus, possibly after tokenization.

        Parameters
        ----------
        word : str
            Word to add to the corpus
        count : int
            Count of word appearances
        doc_count : int
            Count of distinct documents in which word appears


        .. versionadded:: 0.4.0

        """
        if self.transform is not None:
            word = self.transform(word)

        if self.tokenizer is not None:
            self.tokenizer.tokenize(word)
            tokens = self.tokenizer.get_counter()
            for tok in tokens:
                n = tokens[tok] * count
                prior_count, prior_doc_count = self.corpus[tok]
                self.corpus[tok] = (
                    prior_count + n,
                    prior_doc_count + doc_count,
                )
        else:
            prior_count, prior_doc_count = self.corpus[word]
            self.corpus[word] = (
                prior_count + count,
                prior_doc_count + doc_count,
            )

    def gng_importer(self, corpus_file: str) -> None:
        """Fill in self.corpus from a Google NGram corpus file.

        Parameters
        ----------
        corpus_file : file
            The Google NGram file from which to initialize the n-gram corpus


        .. versionadded:: 0.4.0

        """
        with c_open(corpus_file, 'r', encoding='utf-8') as gng:
            for line in gng:
                word, _, count, doc_count = line.rstrip().split('\t')
                if '_' in word:
                    word = word[: word.find('_')]

                self._add_word(word, int(count), int(doc_count))
            self._update_doc_count()

    def idf(self, term: str) -> float:
        r"""Calculate the Inverse Document Frequency of a term in the corpus.

        Parameters
        ----------
        term : str
            The term to calculate the IDF of

        Returns
        -------
        float
            The IDF

        Examples
        --------
        >>> tqbf = 'the quick brown fox jumped over the lazy dog\n\n'
        >>> tqbf += 'and then it slept\n\n and the dog ran off'
        >>> corp = UnigramCorpus(tqbf)
        >>> round(corp.idf('dog'), 10)
        0.6931471806
        >>> round(corp.idf('the'), 10)
        0.6931471806


        .. versionadded:: 0.4.0

        """
        if term in self.corpus:
            count, term_doc_count = self.corpus[term]
            return log1p(self.doc_count / term_doc_count)
        else:
            return float('inf')


if __name__ == '__main__':
    import doctest

    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
