"""Information Content is a corpus-based metrics of synset or sense
specificity.

"""

from collections import Counter
from collections.abc import Callable, Iterable, Iterator
from math import log
from pathlib import Path
from typing import TextIO, TypeAlias

from wn import Synset, Wordnet
from wn._types import AnyPath
from wn.constants import ADJ, ADJ_SAT, ADV, NOUN, VERB
from wn.util import synset_id_formatter

# Just use a subset of all available parts of speech
IC_PARTS_OF_SPEECH = frozenset((NOUN, VERB, ADJ, ADV))
Freq: TypeAlias = dict[str, dict[str | None, float]]


def information_content(synset: Synset, freq: Freq) -> float:
    """Calculate the Information Content value for a synset.

    The information content of a synset is the negative log of the
    synset probability (see :func:`synset_probability`).

    """
    return -log(synset_probability(synset, freq))


def synset_probability(synset: Synset, freq: Freq) -> float:
    """Calculate the synset probability.

    The synset probability is defined as freq(ss)/N where freq(ss) is
    the IC weight for the synset and N is the total IC weight for all
    synsets with the same part of speech.

    Note: this function is not generally used directly, but indirectly
    through :func:`information_content`.

    """
    pos_freq = freq[synset.pos]
    return pos_freq[synset.id] / pos_freq[None]


def _initialize(
    wordnet: Wordnet,
    smoothing: float,
) -> Freq:
    """Populate an Information Content weight mapping to a smoothing value.

    All synsets in *wordnet* are inserted into the dictionary and
    mapped to *smoothing*.

    """
    freq: Freq = {
        pos: {synset.id: smoothing for synset in wordnet.synsets(pos=pos)}
        for pos in IC_PARTS_OF_SPEECH
    }
    # pretend ADJ_SAT is just ADJ
    for synset in wordnet.synsets(pos=ADJ_SAT):
        freq[ADJ][synset.id] = smoothing
    # also initialize totals (when synset is None) for each part-of-speech
    for pos in IC_PARTS_OF_SPEECH:
        freq[pos][None] = smoothing
    return freq


def compute(
    corpus: Iterable[str],
    wordnet: Wordnet,
    distribute_weight: bool = True,
    smoothing: float = 1.0,
) -> Freq:
    """Compute Information Content weights from a corpus.

    Arguments:
        corpus: An iterable of string tokens. This is a flat list of
            words and the order does not matter. Tokens may be single
            words or multiple words separated by a space.

        wordnet: An instantiated :class:`wn.Wordnet` object, used to
            look up synsets from words.

        distribute_weight: If :python:`True`, the counts for a word
            are divided evenly among all synsets for the word.

        smoothing: The initial value given to each synset.

    Example:
        >>> import wn, wn.ic, wn.morphy
        >>> ewn = wn.Wordnet("ewn:2020", lemmatizer=wn.morphy.morphy)
        >>> freq = wn.ic.compute(["Dogs", "run", ".", "Cats", "sleep", "."], ewn)
        >>> dog = ewn.synsets("dog", pos="n")[0]
        >>> cat = ewn.synsets("cat", pos="n")[0]
        >>> frog = ewn.synsets("frog", pos="n")[0]
        >>> freq["n"][dog.id]
        1.125
        >>> freq["n"][cat.id]
        1.1
        >>> freq["n"][frog.id]  # no occurrence; smoothing value only
        1.0
        >>> carnivore = dog.lowest_common_hypernyms(cat)[0]
        >>> freq["n"][carnivore.id]
        1.3250000000000002
    """
    freq = _initialize(wordnet, smoothing)
    counts = Counter(corpus)

    hypernym_cache: dict[Synset, list[Synset]] = {}
    for word, count in counts.items():
        synsets = wordnet.synsets(word)
        num = len(synsets)
        if num == 0:
            continue

        weight = float(count / num if distribute_weight else count)

        for synset in synsets:
            pos = synset.pos
            if pos == ADJ_SAT:
                pos = ADJ
            if pos not in IC_PARTS_OF_SPEECH:
                continue

            freq[pos][None] += weight

            # The following while-loop is equivalent to:
            #
            # freq[pos][synset.id] += weight
            # for path in synset.hypernym_paths():
            #     for ss in path:
            #         freq[pos][ss.id] += weight
            #
            # ...but it caches hypernym lookups for speed

            agenda: list[tuple[Synset, set[Synset]]] = [(synset, set())]
            while agenda:
                ss, seen = agenda.pop()

                # avoid cycles
                if ss in seen:
                    continue

                freq[pos][ss.id] += weight

                if ss not in hypernym_cache:
                    hypernym_cache[ss] = ss.hypernyms()
                agenda.extend((hyp, seen | {ss}) for hyp in hypernym_cache[ss])

    return freq


def load(
    source: AnyPath,
    wordnet: Wordnet,
    get_synset_id: Callable | None = None,
) -> Freq:
    """Load an Information Content mapping from a file.

    Arguments:

        source: A path to an information content weights file.

        wordnet: A :class:`wn.Wordnet` instance with synset
            identifiers matching the offsets in the weights file.

        get_synset_id: A callable that takes a synset offset and part
            of speech and returns a synset ID valid in *wordnet*.

    Raises:

        :class:`wn.Error`: If *wordnet* does not have exactly one
            lexicon.

    Example:

        >>> import wn, wn.ic
        >>> pwn = wn.Wordnet("pwn:3.0")
        >>> path = "~/nltk_data/corpora/wordnet_ic/ic-brown-resnik-add1.dat"
        >>> freq = wn.ic.load(path, pwn)

    """
    source = Path(source).expanduser().resolve(strict=True)
    assert len(wordnet.lexicons()) == 1
    lexid = wordnet.lexicons()[0].id
    if get_synset_id is None:
        get_synset_id = synset_id_formatter(prefix=lexid)

    freq = _initialize(wordnet, 0.0)

    with source.open() as icfile:
        for offset, pos, weight, is_root in _parse_ic_file(icfile):
            ssid = get_synset_id(offset=offset, pos=pos)
            # synset = wordnet.synset(ssid)
            freq[pos][ssid] = weight
            if is_root:
                freq[pos][None] += weight
    return freq


def _parse_ic_file(icfile: TextIO) -> Iterator[tuple[int, str, float, bool]]:
    """Parse the Information Content file.

    A sample of the format is::

        wnver::eOS9lXC6GvMWznF1wkZofDdtbBU
        1740n 1915712 ROOT
        1930n 859272
        2137n 1055337

    """
    next(icfile)  # skip header
    for line in icfile:
        ssinfo, value, *isroot = line.split()
        yield (int(ssinfo[:-1]), ssinfo[-1], float(value), bool(isroot))