#! /usr/bin/python

"""
@author: Thomas Demeester (for more information: thomas.demeester@intec.ugent.be)

This file is part of the TREC FedWeb 2014 Results Merging task.
It aggregates separate nDCG values for the different verticals into the nDCG-IA metric
"""
usage="""
Usage: following input arguments (as location of input files) are required (in this order)
* k:  nDCG-IA@k
* RUN: runfile to be evaluated
* QRELS: qrels file (FW14-QRELS-RM.nodups for the run-specific version with duplicate results removed)
* TRECEVAL: trec_eval executable (e.g., ./trec_eval/trec_eval.9.0/trec_eval)
* RESOURCES: file resources_fedweb2014.txt, containing information on search engines and verticals
* INTENTS: file fedweb14-50test_vr4rm, contain the vertical intent probabilities for each query
* TMP: folder with write access to write dummy qrel files as input to trec_eval
"""

import sys
import subprocess
import os

"""input"""
args = sys.argv
if not len(args)==8 :
    print(usage)
    sys.exit()
k,runfile,qrelsfile,treceval,resourcefile,intentfile,tmp_folder = args[1:]



"""auxiliary variables"""
#map vertical ID to resource ID 
resources = {}
with open(resourcefile,'r') as fID:
    defs = fID.readline()
    for line in fID:
        parts = line.strip().split('\t')
        seID,vID = parts[0],parts[-1]
        if not vID in resources:
            resources[vID] = []
        resources[vID].append(seID)
#create qrels strings, each time setting all resources but from a single vertical to non-relevant. 
qrels_reduced = {}
with open(qrelsfile,'r') as fID:
    qrels = fID.read()
lines = qrels.strip().split('\n')
for vID in resources:    
    qrels_tmp = ''
    for line in lines:
        seID = line.strip().split()[2][:9]
        if not seID in resources[vID]:
            parts = line.strip().split()
            parts[-1] = '0.00'
            line = ' '.join(parts)
        qrels_tmp += line+'\n'   
    qrels_reduced[vID] = qrels_tmp


  
#load vertical intent probabilities
intentprobs = {}
with open(intentfile) as fID:
    for line in fID:
        parts = line.strip().split('\t')
        qID = parts[0]
        intentprobs[qID] = {}
        for s in parts[1:]:
            vID,p = s.split(' ')
            intentprobs[qID][vID] = float(p)
  
"""calculate nDCG-IA"""
all_nDCG = {}
for vID in qrels_reduced:
    #create dummy qrels_reduced file
    qrels_tmp = os.path.join(tmp_folder,qrelsfile.split('/')[-1]+'.%s'%vID)
    with open(qrels_tmp,'w') as fID:
        fID.write(qrels_reduced[vID])
  
    #call trec_eval
    out = subprocess.check_output('%s -q -m ndcg_cut.%s %s %s'%(treceval,k,qrels_tmp,runfile), shell=True)
    for line in out.strip().split('\n')[:-1]:
        [m,qID,v] = line.split()
        if not qID in all_nDCG:
            all_nDCG[qID] = {}
        all_nDCG[qID][vID] = float(v)
    
nDCGIA = {qID:sum([all_nDCG[qID][vID]*intentprobs[qID][vID] for vID in all_nDCG[qID]]) for qID in all_nDCG}
metric = 'ndcg_ia_cut_%s'%k

for qID in sorted(nDCGIA.keys()):
    print("%s\t%s\t%.4f"%(metric,qID,nDCGIA[qID]))    
print("%s\t%s\t%.4f"%(metric,'all',sum(nDCGIA.values())/float(len(nDCGIA))))