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

import utils

try:
    from submitted import baseline, analyze_ambiguous_words
except Exception as e:
    print(e)
    print(traceback.format_exc())


class TestMP03Analysis(unittest.TestCase):

    @weight(10)
    @visibility('visible')
    def test_analyze_ambiguous_words_shape_and_content(self):
        """
        Visible test for the assignment helper:
          analyze_ambiguous_words(train, predicted, gold, k)

        This test checks:
          - function runs without crashing
          - returns k items
          - each item contains required fields
          - returned words are truly ambiguous in TRAIN (seen with >1 tag)
          - distributions are consistent types
        """
        train_set = utils.load_dataset('data/brown-training.txt')
        test_set = utils.load_dataset('data/brown-test.txt')

        # run a tagger to generate predictions (baseline is required anyway)
        predicted = baseline(utils.strip_tags(test_set), train_set)

        # compute ambiguity set from train
        _, multitags_set = utils.get_word_tag_statistics(train_set)

        k = 5
        try:
            result = analyze_ambiguous_words(train_set, predicted, test_set, k=k)
        except Exception as e:
            print(e)
            print(traceback.format_exc())
            self.assertTrue(False, "Error running analyze_ambiguous_words")
            return

        self.assertIsInstance(result, list, "analyze_ambiguous_words must return a list")
        self.assertEqual(len(result), k, f"analyze_ambiguous_words must return exactly k={k} items")

        required_keys = {"word", "train_gold_tag_counts", "test_pred_tag_counts", "test_gold_tag_counts"}

        for item in result:
            # allow either dict or tuple/list, but dict is strongly recommended
            if isinstance(item, dict):
                self.assertTrue(required_keys.issubset(set(item.keys())),
                                f"Each item must include keys {required_keys}")
                word = item["word"]
                train_counts = item["train_gold_tag_counts"]
                pred_counts = item["test_pred_tag_counts"]
                gold_counts = item["test_gold_tag_counts"]
            elif isinstance(item, (tuple, list)) and len(item) >= 4:
                word, train_counts, pred_counts, gold_counts = item[:4]
            else:
                self.assertTrue(False, "Each returned item must be a dict with required keys (recommended), "
                                       "or a tuple/list (word, train_counts, pred_counts, gold_counts).")
                return

            self.assertIsInstance(word, str, "word must be a string")
            self.assertIn(word, multitags_set, f"Returned word '{word}' must be ambiguous in TRAIN (multi-tag)")

            for counts in [train_counts, pred_counts, gold_counts]:
                self.assertIsInstance(counts, dict, "tag count distributions must be dict(tag -> count)")
                for tag, cnt in counts.items():
                    self.assertIsInstance(tag, str, "tag keys must be strings")
                    self.assertIsInstance(cnt, int, "counts must be integers")
                    self.assertGreaterEqual(cnt, 0, "counts must be non-negative")

            # Pred and gold test counts should sum to the same number of occurrences of the word in the test set
            pred_total = sum(pred_counts.values())
            gold_total = sum(gold_counts.values())
            self.assertEqual(pred_total, gold_total,
                             f"Pred vs gold test totals must match for word '{word}' (same word occurrences).")
