"""
norvig.nlp – AIMA Chapter 22
==============================
Uses NLTK as the NLP engine throughout.

    from ml_ai_library.norvig.nlp import (
        NGramModel,      # wraps nltk.ngrams + nltk.FreqDist
        CYKParser,       # wraps nltk.ChartParser (CYK strategy)
        viterbi_pos,     # nltk.HiddenMarkovModelTrainer
        make_grammar,    # nltk.CFG.fromstring helper
    )

AIMA chapter map:
  NGramModel   → Ch 22.1 (n-gram language models)
  CYKParser    → Ch 22.3 (CYK chart parsing)
  viterbi_pos  → Ch 22.4 (HMM POS tagging)
"""

from __future__ import annotations
import math
from typing import Any, Dict, List, Optional, Sequence, Tuple

# NLTK is the engine for all NLP functionality --------------------------------
import nltk
from nltk import FreqDist, ngrams as _nltk_ngrams
from nltk.grammar import CFG, Nonterminal
from nltk.parse.chart import (
    ChartParser,
    BU_LC_STRATEGY,            # bottom-up left-corner chart = close to CYK
)

# Ensure punkt tokenizer data is available for demos
try:
    nltk.data.find("tokenizers/punkt")
except LookupError:
    nltk.download("punkt", quiet=True)


# ---------------------------------------------------------------------------
# N-Gram Language Model (AIMA §22.1)
# ---------------------------------------------------------------------------

class NGramModel:
    """
    N-gram language model backed by nltk.ngrams + nltk.FreqDist.

    Parameters
    ----------
    n       : order (1=unigram, 2=bigram, 3=trigram …)
    laplace : Laplace (add-1) smoothing count
    """

    def __init__(self, n: int = 2, laplace: float = 1.0):
        self.n       = n
        self.laplace = laplace
        self._hist: Dict[tuple, FreqDist] = {}   # context -> FreqDist(next_word)
        self._vocab: set = set()

    def train(self, tokens: Sequence[str]):
        """Train on a flat list of word tokens.  Appends to existing model."""
        self._vocab.update(tokens)
        # Pad with special start/end tokens
        padded = list(tokens)
        for gram in _nltk_ngrams(padded, self.n, pad_left=True,
                                 pad_right=True, left_pad_symbol="<s>",
                                 right_pad_symbol="</s>"):
            context, word = gram[:-1], gram[-1]
            if context not in self._hist:
                self._hist[context] = FreqDist()
            self._hist[context][word] += 1

    def log_prob(self, word: str, context: tuple) -> float:
        """Log₂ probability of `word` given `context` (Laplace-smoothed)."""
        V    = len(self._vocab) + 1
        dist = self._hist.get(context, FreqDist())
        cnt  = dist[word] + self.laplace
        tot  = dist.N() + self.laplace * V
        return math.log2(cnt / tot)

    def perplexity(self, tokens: Sequence[str]) -> float:
        """Perplexity of token sequence under the model."""
        total_log = 0.0
        count = 0
        padded = list(tokens)
        for gram in _nltk_ngrams(padded, self.n, pad_left=True,
                                 pad_right=False, left_pad_symbol="<s>"):
            context, word = gram[:-1], gram[-1]
            total_log += self.log_prob(word, context)
            count += 1
        return 2 ** (- total_log / max(count, 1))

    def top_next(self, context: tuple, k: int = 5) -> List[Tuple[str, float]]:
        """Return the top-k most probable next words given `context`."""
        dist = self._hist.get(context, FreqDist())
        cands = list(self._vocab or dist.keys())
        scored = [(w, self.log_prob(w, context)) for w in cands]
        return sorted(scored, key=lambda x: -x[1])[:k]


# ---------------------------------------------------------------------------
# CYK Parser via NLTK ChartParser (AIMA §22.3)
# ---------------------------------------------------------------------------

def make_grammar(rules_string: str) -> CFG:
    """
    Build a nltk.CFG from a multi-line string.

    Example
    -------
    >>> g = make_grammar('''
    ...   S  -> NP VP
    ...   NP -> Det N
    ...   VP -> V NP
    ...   Det -> 'the' | 'a'
    ...   N   -> 'cat' | 'dog'
    ...   V   -> 'sees' | 'chases'
    ... ''')
    """
    return CFG.fromstring(rules_string)


class CYKParser:
    """
    CYK-style chart parser backed by nltk.ChartParser with the
    BU_LC_STRATEGY (bottom-up, left-corner filtering) strategy.

    Parameters
    ----------
    grammar : CFG or a rules string (passed to make_grammar)
    """

    def __init__(self, grammar):
        if isinstance(grammar, str):
            grammar = CFG.fromstring(grammar)
        self._parser = ChartParser(grammar, strategy=BU_LC_STRATEGY)

    def parse(self, tokens: List[str]):
        """
        Parse `tokens`.  Returns list of nltk.Tree objects.
        Each tree has a string representation and a pretty-print method.

        Example
        -------
        >>> parser = CYKParser(make_grammar("S -> 'a'"))
        >>> trees = parser.parse(['a'])
        >>> str(trees[0])
        "(S 'a')"
        """
        return list(self._parser.parse(tokens))

    def parse_str(self, sentence: str) -> List:
        """Tokenise `sentence` on whitespace then parse."""
        return self.parse(sentence.strip().split())


# ---------------------------------------------------------------------------
# Viterbi POS Tagging via NLTK HMM (AIMA §22.4)
# ---------------------------------------------------------------------------

def viterbi_pos(train_sents: Optional[List] = None):
    """
    Train and return an NLTK HiddenMarkovModelTagger (Viterbi decoding).

    Parameters
    ----------
    train_sents : list of list of (word, tag) tuples.
                  If None, downloads and uses the Brown corpus (tiny sample).

    Returns
    -------
    tagger : HiddenMarkovModelTagger with a .tag(tokens) method.

    Example
    -------
    >>> tagger = viterbi_pos()
    >>> tagger.tag(['The', 'cat', 'runs'])
    [('The', 'AT'), ('cat', 'NN'), ('runs', 'VBZ')]
    """
    from nltk.tag import hmm as _hmm
    if train_sents is None:
        try:
            nltk.data.find("corpora/brown")
        except LookupError:
            nltk.download("brown", quiet=True)
        from nltk.corpus import brown
        train_sents = brown.tagged_sents()[:2000]

    trainer = _hmm.HiddenMarkovModelTrainer()
    return trainer.train_supervised(train_sents)
