#!/usr/bin/python
# refines the query: deduplication of queries
# Usage: python refine_query.py <old-query-file> <old-qrel-file> 
# Input: old-query-file old-query-qrel
# Output: new-query.txt new-query.qrel

import sys
import os
import glob
from xml.etree import ElementTree as ET
import string
from collections import Counter
import re, math

class QueryList(object):
    def __init__(self, p_num, p_title, p_abstract, q_num, text, rel_cid):
        self.p_num = p_num
        self.p_title = p_title
        self.p_abstract = p_abstract
        self.q_num = q_num
        self.text = text
        self.rel_cid = rel_cid

    def __repr__(self):
        return '{}: {} {}'.format(self.__class__.__name__,
                                  self.p_num,
                                  self.p_title,
                                  self.p_abstract,
                                  self.q_num,
                                  self.text,
                                  self.rel_cid)

query_qrel_cid_list=[]

def is_empty(any_structure):
    if any_structure:
        # print('Structure is not empty.')
        return False
    else:
        # print('Structure is empty.')
        return True
""" Similarity measurement"""
WORD = re.compile(r'\w+')


def binary_search(a, x, lo=0, hi=None):
    if hi is None:
        hi = len(a)
    while lo < hi:
        mid = (lo+hi)//2
        midval = a[mid]
        # print "Comparing <%s><%s>"%(a[mid], x)
        if midval < x:
            lo = mid+1
        elif midval > x: 
            hi = mid
        else:
            return mid
    return -1

def query_text_compare(x, y):
   if x.text > y.text:
      return 1
   elif x.text == y.text:
      return 0
   else:  #x.text < y.text
      return -1

def get_text(item):
    return item.text


def parse_query(query, qrel):
    rootElement = ET.parse(query).getroot()
    query_txt=open("new_query.txt","w")
    query_qrel=open("new_query.qrel","w")
    print >>query_txt, "<query>"
    for sub in rootElement:
        print >>query_txt, "<top>"
        # initialize(sub)
        # print "Tag: ",sub.tag
        paper_num=""
        paper_title=""
        paper_abstract=""
        query_num=[]
        text=[]
        relevant_cid=[]
        query_num_count=0
        for subsub in sub:
            # print "SubTag: ",subsub.tag
            if subsub.tag == "paper_num":
                paper_num = subsub.text
                print >>query_txt, "<paper_num> %d </paper_num>" % (int(paper_num))
                query_num_count = 0
            if subsub.tag == "paper_title":
                paper_title = subsub.text
                print >>query_txt, "<paper_title> %s </paper_title>" % (paper_title)
            if subsub.tag == "paper_abstract":
                paper_abstract = subsub.text
                print >>query_txt, "<paper_abstract> %s </paper_abstract>" % (paper_abstract.encode('utf8'))
            if subsub.tag == "query_num":
                query_num.append(subsub.text)
                # print "qno: ", subsub.text
                ## linear search for the relevant paper cid for the 'query_num' 
                for q in query_qrel_cid_list:
                    # print ("<%s><%s>"%(q[0],subsub.text.strip()))
                    if q[0] == subsub.text.strip():
                        # print "found", q[1], q[0]
                        relevant_cid.append(q[1])
                        break
            if subsub.tag == "text":
                # print "text"
                text.append(subsub.text)

        # print len(text)

        query_list=[]
        for i in xrange(0,len(text)):
            # print i
            # print query_num[i]
            # print text[i].encode('utf8')
            query_list.append(QueryList(paper_num, paper_title, paper_abstract, query_num[i], text[i], relevant_cid[i]))


        # for i in xrange(0, len(query_list)):
        #     print query_list[i].text.encode('utf8')
        query_list_sorted= sorted(query_list, key=get_text)
        last_text=""
        last_rel_cid=""
        final_query=[]
        final_qrel=[]
        for i in xrange(0, len(query_list_sorted)):
            if last_text == query_list_sorted[i].text:
                # print query_list_sorted[i].rel_cid
                print >>query_qrel, ("%d 0 %d 1") %(int(str_context_count), int(query_list_sorted[i].rel_cid))
            else:
                # print query_list_sorted[i].q_num, query_list_sorted[i].text.encode('utf8')
                last_text = query_list_sorted[i].text
                last_rel_cid = query_list_sorted[i].rel_cid
                query_num_count = query_num_count + 1
                if query_num_count < 10:
                    str_context_count = "%d0%d" % (int(paper_num), query_num_count)
                    print >>query_txt, "<query_num> %s </query_num>" % str_context_count #query_list_sorted[i].rel_cid
                if query_num_count >= 10 and query_num_count <99:
                    str_context_count = "%s%d" % (int(paper_num), query_num_count)
                    print >>query_txt, "<query_num> %s </query_num>" % str_context_count #query_list_sorted[i].rel_cid
                print >>query_txt, "<text> %s </text>" % query_list_sorted[i].text.encode('utf8')
                print >>query_qrel, ("%d 0 %d 1") %(int(str_context_count), int(query_list_sorted[i].rel_cid))

#                print >>query_txt, query_list_sorted[i].rel_cid
    print >>query_txt, "<query>"

def parse_qrel(qrel_file):
    with open(qrel_file, "r") as qrel:
        for line in qrel:
            words = line.split(" ")
            # print words[0], words[2], words[3]
            query_qrel_cid_list.append([words[0], words[2]])
            # array.append(line)


if __name__ == "__main__":
    if len(sys.argv) != 1:
        parse_qrel(sys.argv[2])
        parse_query(sys.argv[1], sys.argv[2])
    else:
        print "Usage: python refine_query.py <old-query-file> <old-qrel-file>"

