from collections import Counter
import random
import unittest
from example import SearchIndex, tokenize


class SearchTests(unittest.TestCase):
    def test_normalization_and_and_query(self):
        index = SearchIndex()
        index.upsert('a', 'Ｐｙｔｈｏｎ 图算法 Cafe\u0301')
        self.assertEqual(index.search('python 图'), [('a', 2)])
        self.assertEqual(index.search('CAFÉ'), [('a', 1)])
        self.assertEqual(index.search('missing'), [])
        self.assertEqual(index.search(''), [])

    def test_replace_remove_and_tie(self):
        index = SearchIndex()
        index.upsert('b', 'graph')
        index.upsert('a', 'graph')
        self.assertEqual(index.search('graph graph'), [('a', 1), ('b', 1)])
        index.upsert('a', 'tree')
        self.assertEqual(index.search('graph'), [('b', 1)])
        self.assertTrue(index.remove('b'))
        self.assertFalse(index.remove('b'))
        self.assertEqual(index.search('graph'), [])

    def test_against_scan(self):
        rng = random.Random(21)
        index = SearchIndex()
        texts = {str(i): ' '.join(rng.choice(['图', '算法', 'python', 'tree'])
                                 for _ in range(12)) for i in range(20)}
        for doc_id, text in texts.items():
            index.upsert(doc_id, text)
        for query in ('图 算法', 'python', 'tree python', '图', 'missing'):
            terms = set(tokenize(query))
            scored = []
            for doc_id, text in texts.items():
                counts = Counter(tokenize(text))
                if terms and all(counts[term] for term in terms):
                    scored.append((-sum(counts[term] for term in terms), doc_id))
            expected = [(doc_id, -negative) for negative, doc_id in sorted(scored)[:5]]
            self.assertEqual(index.search(query), expected)


if __name__ == '__main__':
    unittest.main()