"""
ECE 448 / CS 440 MP03: HMM POS Tagging (Brown Corpus)

STUDENT DISTRIBUTION TEMPLATE

You will submit this file on Gradescope.

Required:
  - baseline(test, train)
  - viterbi(test, train)


IMPORTANT DESIGN NOTE:
You are expected to implement several helper functions below. The autograder may test
these helpers directly to pinpoint issues.

Dataset format:
- train: list of sentences; each sentence is a list of (word, tag) pairs
- test:  list of sentences; each sentence is a list of words
The provided utils.load_dataset() typically includes special tokens "START" and "END"
as words with tags "START" and "END".

Rules:
- Use only the Python standard library (collections, math, etc.) and numpy if desired.
- Do NOT use NLP libraries like nltk/spacy.
"""

from collections import defaultdict, Counter
import math

# Smoothing constant (Laplace/add-epsilon). You may tune for better performance.
EPS = 1e-5


# ---------------------------------------------------------------------
# Part 1: Baseline tagger helpers
# ---------------------------------------------------------------------

def compute_word_tag_counts(train):
    """
    Compute counts needed for the baseline tagger.

    Args:
        train: list of sentences, each a list of (word, tag) tuples
        
    Returns:
        tuple: (word_tag_counts, tag_counts) where:
            - word_tag_counts: dict[word] -> Counter(tag -> count)
            - tag_counts: Counter(tag -> count) over the entire training set
            
    Example:
        >>> train = [[('the', 'DT'), ('dog', 'NN')], [('a', 'DT'), ('cat', 'NN')]]
        >>> word_tag_counts, tag_counts = compute_word_tag_counts(train)
        >>> word_tag_counts['the']['DT']
        1
        >>> tag_counts['NN']
        2
    """
    raise NotImplementedError


def get_most_frequent_tag(tag_counts):
    """
    Return the most common tag overall. Used as a fallback tag for unseen words.
    
    Args:
        tag_counts: Counter(tag -> count) over the entire training set
        
    Returns:
        str: The tag with the highest count in tag_counts
        
    Example:
        >>> tag_counts = Counter({'NN': 100, 'VB': 50, 'DT': 75})
        >>> get_most_frequent_tag(tag_counts)
        'NN'
    """
    raise NotImplementedError


def predict_baseline_tag(word, word_tag_counts, fallback_tag):
    """
    For a given word, return the most frequent tag seen with that word in training.
    If unseen, return fallback_tag.
    
    Args:
        word: str, the word to predict a tag for
        word_tag_counts: dict[word] -> Counter(tag -> count)
        fallback_tag: str, tag to return if word is not in word_tag_counts
        
    Returns:
        str: The most frequent tag for the word, or fallback_tag if unseen
        
    Example:
        >>> word_tag_counts = {'dog': Counter({'NN': 10, 'VB': 2}), 'cat': Counter({'NN': 8})}
        >>> predict_baseline_tag('dog', word_tag_counts, 'NN')
        'NN'
        >>> predict_baseline_tag('bird', word_tag_counts, 'NN')
        'NN'
    """
    raise NotImplementedError


def baseline(test, train):
    """
    Baseline tagger:
    - For each word, predict the most frequent tag associated with that word in TRAIN.
    - For unseen words, predict the globally most frequent tag in TRAIN.

    Args:
        test: list of sentences, each a list of words
        train: list of sentences, each a list of (word, tag) pairs

    Returns:
        predicted: list of sentences, each a list of (word, tag) pairs
        
    Note: Use the helper functions compute_word_tag_counts(), get_most_frequent_tag(), 
          and predict_baseline_tag() to implement this function.
    """
    # Hint: use compute_word_tag_counts(), get_most_frequent_tag(), predict_baseline_tag()
    raise NotImplementedError


# ---------------------------------------------------------------------
# Part 2: HMM training helpers (counts + smoothed log-probabilities)
# ---------------------------------------------------------------------

def collect_hmm_counts(train):
    """
    Collect HMM counts from the training data.

    Args:
        train: list of sentences, each a list of (word, tag) tuples
        
    Returns:
        tuple: (tag_counts, trans_counts, emit_counts, tagset, vocab) where:
            - tag_counts: Counter(tag -> count) - number of times each tag appears
            - trans_counts: dict[prev_tag] -> Counter(next_tag -> count) - transition counts
            - emit_counts: dict[tag] -> Counter(word -> count) - emission counts  
            - tagset: list of unique tags, sorted alphabetically
            - vocab: list of unique words (vocabulary), sorted alphabetically
            
    Note: 
        - For transitions, consider consecutive tags within each sentence
        - tagset and vocab should be deterministic (sorted) for consistent testing
        
    Example:
        >>> train = [[('START', 'START'), ('the', 'DT'), ('dog', 'NN'), ('END', 'END')]]
        >>> tag_counts, trans_counts, emit_counts, tagset, vocab = collect_hmm_counts(train)
        >>> trans_counts['START']['DT']
        1
        >>> emit_counts['NN']['dog'] 
        1
    """
    raise NotImplementedError


def compute_transition_logprobs(trans_counts, tagset, smoothing=EPS):
    """
    Compute smoothed log transition probabilities log P(tag | prev_tag).

    Args:
        trans_counts: dict[prev_tag] -> Counter(next_tag -> count)
        tagset: list of unique tags
        smoothing: float, smoothing constant (default EPS)
        
    Returns:
        dict: trans_logprob[(prev_tag, tag)] -> float (log probability)
        
    Note:
        - Use add-epsilon (Laplace) smoothing: (count + epsilon) / (total + epsilon * |tagset|)
        - Return natural logarithm of probabilities
        - Should include entries for all (prev_tag, tag) pairs in tagset × tagset
        
    Example:
        >>> trans_counts = {'START': Counter({'DT': 1})}
        >>> tagset = ['START', 'DT', 'END'] 
        >>> logprobs = compute_transition_logprobs(trans_counts, tagset, 0.1)
        >>> ('START', 'DT') in logprobs
        True
    """
    raise NotImplementedError


def compute_emission_logprobs(emit_counts, tag_counts, vocab, smoothing=EPS):
    """
    Compute smoothed log emission probabilities log P(word | tag).

    Args:
        emit_counts: dict[tag] -> Counter(word -> count)
        tag_counts: Counter(tag -> count) 
        vocab: list of words in vocabulary
        smoothing: float, smoothing constant (default EPS)
        
    Returns:
        tuple: (emit_logprob, unk_logprob) where:
            - emit_logprob: dict[(tag, word)] -> float (log probability)
            - unk_logprob: dict[tag] -> float (log prob for unknown words) OR None
            
    Note:
        You must handle UNKNOWN words at test time. Two approaches:
        
        (A) Add explicit '<UNK>' entries: emit_logprob[(tag, '<UNK>')] and return unk_logprob=None
        (B) Return unk_logprob[tag] for each tag and handle unknowns in viterbi_decode()
        
        Use add-epsilon smoothing: (count + epsilon) / (total + epsilon * |vocab|)
        For approach A, |vocab| includes the '<UNK>' token.
        
    Example:
        >>> emit_counts = {'NN': Counter({'dog': 2, 'cat': 1})}
        >>> tag_counts = Counter({'NN': 3})
        >>> vocab = ['dog', 'cat']
        >>> emit_logprob, unk_logprob = compute_emission_logprobs(emit_counts, tag_counts, vocab)
        >>> ('NN', 'dog') in emit_logprob or unk_logprob is not None
        True
    """
    raise NotImplementedError


# ---------------------------------------------------------------------
# Part 3: Viterbi decoding helpers
# ---------------------------------------------------------------------

def viterbi_decode(sentence, tagset, trans_logprob, emit_logprob, unk_logprob=None):
    """
    Run Viterbi decoding on ONE sentence.

    Args:
        sentence: list of words (typically includes START and END tokens)
        tagset: list of tags (typically includes START and END tags)
        trans_logprob: dict[(prev_tag, tag)] -> float (log P(tag | prev_tag))
        emit_logprob: dict[(tag, word)] -> float (log P(word | tag))
        unk_logprob: dict[tag] -> float (log P(UNK | tag)) if using approach B, else None

    Returns:
        list: best_tags, same length as sentence
        
    Note:
        - Implement standard Viterbi algorithm with dynamic programming
        - Handle unknown words using unk_logprob or '<UNK>' entries in emit_logprob
        - Return the most likely tag sequence for the given sentence
        - Use log probabilities throughout to avoid numerical underflow
        
    Example:
        >>> sentence = ['START', 'the', 'dog', 'END']
        >>> tagset = ['START', 'DT', 'NN', 'END']
        >>> # ... (with appropriate trans_logprob and emit_logprob)
        >>> tags = viterbi_decode(sentence, tagset, trans_logprob, emit_logprob, unk_logprob)
        >>> len(tags) == len(sentence)
        True
    """
    raise NotImplementedError


def viterbi(test, train):
    """
    Full HMM tagger:
      1) collect counts from training data
      2) build smoothed log-probability tables  
      3) Viterbi decode each test sentence
      4) return (word, tag) pairs

    Args:
        test: list of sentences (each sentence is a list of words)
        train: list of sentences (each sentence is a list of (word, tag) pairs)

    Returns:
        predicted: list of sentences, each a list of (word, tag) pairs
        
    Note: Use the helper functions collect_hmm_counts(), compute_transition_logprobs(),
          compute_emission_logprobs(), and viterbi_decode() to implement this function.
    """
    raise NotImplementedError



def analyze_ambiguous_words(train_sentences, predicted_sentences, tag_sentences, k=5):
    """
    (Write-up helper) Identify the top-k ambiguous words that your model mis-tags most on TEST.

    Definitions:
      - A word is *ambiguous* if it appears with >1 tag in TRAIN.
      - A word is *mis-tagged* at a test position if predicted_tag != gold_tag.

    Inputs:
      train_sentences: list of tagged training sentences (list of (word, tag))
      predicted_sentences: model outputs on test (list of (word, predicted_tag))
      tag_sentences: gold-labeled test sentences (list of (word, gold_tag))
      k: number of words to return

    Output:
      A list of up to k dicts; each dict has keys:
        - "word": str
        - "train_gold_tag_counts": dict(tag -> count) in TRAIN
        - "test_pred_tag_counts": dict(tag -> count) predicted on TEST
        - "test_gold_tag_counts": dict(tag -> count) gold on TEST

    Ranking / tie-breaking (deterministic):
      1) descending mistakes[w]
      2) if tie: descending test frequency
      3) if tie: lexicographic word (ascending)

    If fewer than k ambiguous words have mistakes>0:
      - include all with mistakes>0, then fill with ambiguous words that appear on TEST with mistakes==0,
        ranked by test frequency then lexicographic.
      - if still fewer than k exist, return as many as exist.
    """
    raise NotImplementedError
