#!/usr/bin/env python3

import subprocess
import tempfile
import argparse
import os
import gzip

argparser = argparse.ArgumentParser(
    description='Evaluate a TREC 2021 Podcast segment retrieval run using trec_eval',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)

argparser.add_argument('--treceval',
                       help='Path to trec_eval executable',
                       default='/retrieval/evaluations/trec/trec30/scripts.process/trec_eval_v9.macos')
argparser.add_argument('--qrels',
                       help='Qrels file',
                       default='/retrieval/evaluations/trec/trec30/podcast/eval/qrels.segments')
argparser.add_argument('-m', '--messy', action='store_true',
                       help='Keep temporary files')
argparser.add_argument('runfile', help='Path to run to evaluate')

args = argparser.parse_args()

def open_either(filename):
    if filename.endswith('.gz'):
        return gzip.open(filename, 'r')
    else:
        return open(filename, 'r')

rankings = {}
qrels = {}

runtag = None

# Split the four rankings into four separate runfiles
with open_either(args.runfile) as runfile:
    for line in runfile:
        topic, which, docid, rank, sim, runtag = line.split()
        if which not in rankings:
            rankings[which] = tempfile.NamedTemporaryFile(prefix=f'run.{which}.',
                                                          mode='w+',
                                                          delete=False)
        print(line, file=rankings[which])

for which, ranking in rankings.items():
    rankings[which] = ranking.name
    ranking.flush()
    ranking.close()
    qrels[which] = tempfile.NamedTemporaryFile(prefix=f'qrels.{which}.',
                                               mode='w+',
                                               delete=False)
    
# construct qrels files for each ranking condition
with open(args.qrels, 'r') as qrels_file:
    for line in qrels_file:
        if line.startswith('#'):
            continue
        topic, _, docno, rel, ent, subj, disc = line.split()
        rel = int(rel)
        ent = int(ent)
        subj = int(subj)
        disc = int(subj)
        for which in qrels:
            if rel == 0:
                print(topic, 0, docno, 0, file=qrels[which])
            else: 
                gain = rel
                if (which == 'QR') or \
                   (which == 'QE' and ent > 0) or \
                   (which == 'QS' and subj > 0) or \
                   (which == 'QD' and disc > 0):
                    print(topic, 0, docno, gain, file=qrels[which])

# run trec_eval using each runfile and qrels file
for which, qrelsfile in qrels.items():
    qrels[which] = qrelsfile.name
    qrelsfile.flush()
    qrelsfile.close()

    output_filename = f'{runtag}.{which}.eval'
    with open(output_filename, 'w') as eval_output:
        result = subprocess.run([args.treceval,
                                 '-M 1000',
                                 '-q',
                                 '-mall_trec',
                                 qrels[which],
                                 rankings[which]],
                                stdout=eval_output)
        if result.returncode != 0:
            raise OSError(result.returncode, result.stderr)

    if args.messy:
        print(which, qrels[which], rankings[which])
    else:
        os.unlink(qrels[which])
        os.unlink(rankings[which])