from typing import Dict
from typing import List
import collections


def load_run(path: str) -> Dict[str, List[str]]:
    """Loads run into a dict of key: query_id, value: list of candidate doc ids."""
    run = collections.OrderedDict()
    with open(path) as f:
        for line in f:
            query_id, _, doc_title, rank, _, _ = line.split()
            if query_id not in run:
                run[query_id] = []
            run[query_id].append((doc_title, int(rank)))

    # Sort candidate docs by rank.
    sorted_run = collections.OrderedDict()
    for query_id, doc_titles_ranks in run.items():
        doc_titles_ranks.sort(key=lambda x: x[1])
        doc_titles = [doc_titles for doc_titles, _ in doc_titles_ranks]
        sorted_run[query_id] = doc_titles

    return sorted_run