from collections import Counter, defaultdict
import heapq
import unicodedata


def tokenize(text):
    text = unicodedata.normalize('NFKC', text).casefold()
    tokens, word = [], []
    for char in text:
        if '\u4e00' <= char <= '\u9fff':
            if word:
                tokens.append(''.join(word))
                word.clear()
            tokens.append(char)
        elif char.isalnum():
            word.append(char)
        elif word:
            tokens.append(''.join(word))
            word.clear()
    if word:
        tokens.append(''.join(word))
    return tokens


class SearchIndex:
    def __init__(self):
        self.documents = {}
        self.postings = defaultdict(set)

    def remove(self, doc_id):
        old = self.documents.pop(doc_id, None)
        if old is None:
            return False
        for token in old:
            self.postings[token].remove(doc_id)
            if not self.postings[token]:
                del self.postings[token]
        return True

    def upsert(self, doc_id, text):
        if not isinstance(doc_id, str) or not isinstance(text, str):
            raise ValueError('document ID and text must be strings')
        self.remove(doc_id)
        counts = Counter(tokenize(text))
        self.documents[doc_id] = counts
        for token in counts:
            self.postings[token].add(doc_id)

    def search(self, query, limit=5):
        if limit < 0:
            raise ValueError('limit must be non-negative')
        terms = set(tokenize(query))
        if not terms or not limit or any(term not in self.postings for term in terms):
            return []
        groups = sorted((self.postings[term] for term in terms), key=len)
        candidates = groups[0].copy()
        for group in groups[1:]:
            candidates.intersection_update(group)
        scored = ((-sum(self.documents[doc_id][term] for term in terms), doc_id)
                  for doc_id in candidates)
        return [(doc_id, -negative) for negative, doc_id in heapq.nsmallest(limit, scored)]


if __name__ == '__main__':
    import argparse
    import json
    from pathlib import Path
    parser = argparse.ArgumentParser(description='Local AND search')
    parser.add_argument('query', nargs='?', default='图 算法')
    parser.add_argument('corpus', nargs='?')
    args = parser.parse_args()
    corpus = {'a': '图算法与最短路径', 'b': 'Python 算法 算法', 'c': '图算法 图算法'}
    if args.corpus:
        corpus = json.loads(Path(args.corpus).read_text(encoding='utf-8'))
    if not isinstance(corpus, dict):
        parser.error('corpus must be a JSON object mapping IDs to text')
    index = SearchIndex()
    for doc_id, text in corpus.items():
        index.upsert(doc_id, text)
    print(index.search(args.query))