# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['calculate_score', 'do_run']

# Cell

def calculate_score(fn_qrels='data/robust/qrels.robust2004.txt', prediction="score.txt"):
    cmd = "eval/trec_eval.9.0.4/trec_eval {judgement} {output} -m map -m recip_rank".format(output=prediction, judgement=fn_qrels)
    pargs = shlex.split(cmd)
    p = subprocess.Popen(pargs, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    pout, perr = p.communicate()
    print("running {}".format(cmd))
    if sys.version_info[0] < 3:
        lines = pout.split('\n')
    else:
        lines = pout.split(b'\n')
    MAP = float(lines[0].strip().split()[-1])
    MRR = float(lines[1].strip().split()[-1])
    print("MAP: {}".format(MAP))
    print("MRR: {}".format(MRR))
    return MAP, MRR

# Cell

def do_run(file, topics, searcher):
    with open(file, 'w') as runfile:
        cnt = 0
        print('Running {} queries in total'.format(len(topics)))
        start = timer()
        for id in tqdm(topics):
            query = topics[id]['title'].encode('utf-8');
            # see https://github.com/kivy/pyjnius/issues/437 on why we need to encode
            hits = searcher.search(query, 1000)
            for i in range(0, len(hits)):
                _ = runfile.write('{} Q0 {} {} {:.6f} Anserini\n'.format(id, hits[i].docid, i+1, hits[i].score))
        end = timer()
        timeneeded = end - start
        with open('time-' + file, 'w') as timefile:
            _ = timefile.write('{}'.format(timeneeded))
        print(end-start)