import math
import unittest
import traceback
from gradescope_utils.autograder_utils.decorators import weight, visibility

import utils

try:
    import submitted
except Exception as e:
    print(e)
    print(traceback.format_exc())


class TestMP03Helpers(unittest.TestCase):

    @weight(5)
    @visibility('visible')
    def test_compute_word_tag_counts(self):
        """Sanity check: baseline count tables have expected types and non-empty content."""
        train_set = utils.load_dataset('data/synthetic_training.txt')
        wtc, tc = submitted.compute_word_tag_counts(train_set)

        self.assertIsInstance(wtc, dict)
        self.assertIsInstance(tc, dict.__mro__[1])  # Counter is a dict subclass
        self.assertGreater(len(tc), 0)
        # A common bug: forgetting to count tags at all
        self.assertGreater(sum(tc.values()), 0)

    @weight(5)
    @visibility('visible')
    def test_hmm_counts_and_logprobs_normalize(self):
        """
        Sanity check HMM helpers:
        - collect_hmm_counts returns non-empty tagset/vocab
        - transition rows approximately normalize (in probability space)
        """
        train_set = utils.load_dataset('data/synthetic_training.txt')
        tag_counts, trans_counts, emit_counts, tagset, vocab = submitted.collect_hmm_counts(train_set)

        self.assertIsInstance(tagset, list)
        self.assertGreater(len(tagset), 0)
        self.assertIsInstance(vocab, list)
        self.assertGreater(len(vocab), 0)

        trans_lp = submitted.compute_transition_logprobs(trans_counts, tagset)

        # Pick one prev_tag that appears in transitions
        prev_candidates = [t for t in tagset if t in trans_counts and sum(trans_counts[t].values()) > 0]
        self.assertGreater(len(prev_candidates), 0)
        prev = prev_candidates[0]

        # Check normalization: sum_t exp(logP(t|prev)) ≈ 1
        # Allow loose tolerance due to smoothing/float
        s = 0.0
        for t in tagset:
            key = (prev, t)
            self.assertIn(key, trans_lp, "Missing transition entry for (prev_tag, tag).")
            s += math.exp(trans_lp[key])
        self.assertTrue(0.95 <= s <= 1.05, f"Transition probabilities for prev_tag={prev} do not normalize (~1). Got {s}.")
