#
# msmarco doc: create the train.tsv triples  
# -------------------------------

import random
random.seed(42)

import argparse
import os
import sys
from tqdm import tqdm
sys.path.append(os.getcwd())

from matchmaker.evaluation.msmarco_eval import *
from collections import defaultdict
from matchmaker.dataloaders.bling_fire_tokenizer import BlingFireTokenizer

#
# config
#
parser = argparse.ArgumentParser()

parser.add_argument('--out-file', action='store', dest='out_file',
                    help='training output text file location', required=True)

parser.add_argument('--out-file-ids', action='store', dest='out_file_ids',
                    help='training output ids file location', required=True)

parser.add_argument('--candidate-file', action='store', dest='candidate_file',
                    help='trec ranking file location (lucene output)', required=True)

parser.add_argument('--collection-file', action='store', dest='collection_file',
                    help='collection.tsv location', required=True)

parser.add_argument('--query-file', action='store', dest='query_file',
                    help='query.tsv location', required=True)

parser.add_argument('--qrel', action='store', dest='qrel_file',
                    help='qrel location', required=True)


args = parser.parse_args()

max_triples = 10_000_000
max_doc_char_length = 150_000

max_doc_token_length = 10000

#
# load data 
# -------------------------------
# 
collection = {}
#collection_length = {} 
tokenizer = BlingFireTokenizer()
with open(args.collection_file,"r",encoding="utf8") as collection_file:
    for line in tqdm(collection_file):
        ls = line.split("\t") # id<\t>text ....
        _id = ls[0]
        max_char_doc = ls[1].rstrip()[:max_doc_char_length]
        collection[_id] = max_char_doc
        #collection_length[_id] = len(tokenizer.tokenize(max_char_doc))

queries = {}
with open(args.query_file,"r",encoding="utf8") as query_file:
    for line in tqdm(query_file):
        ls = line.split("\t") # id<\t>text ....
        _id = ls[0]
        queries[_id] = ls[1].rstrip()

qrels = load_reference(args.qrel_file)


#
# produce output
# -------------------------------
#  

triples = []

stats = defaultdict(int)

with open(args.candidate_file,"r",encoding="utf8") as candidate_file:

    for line in tqdm(candidate_file):
        #if random.random() <= 0.5: continue #skip some entries for faster processing
        [topicid, _ , unjudged_docid, rank, _ , _ ] = line.split()

        #if int(rank) <= 100:
        #    #if random.random() < 0.7: continue # skip 70% of candidates to speed up things...
        #    #else:
        #    stats['< 100 sampling count'] += 1
        #else:
        #    if random.random() <= 0.9: continue # skip 90% of candidates assumong top1k -> same number of samples from 0-100 as 101 - 1000
        #    else:
        #        stats['> 100 sampling count'] += 1

        if topicid not in queries or topicid not in qrels: # added: because we carved out the validation qrels from the train -> so there are some missing
            stats['skipped'] += 1
            continue

        #assert topicid in qrels
        assert unjudged_docid in collection

        # Use topicid to get our positive_docid
        positive_docid = random.choice(qrels[topicid])
        assert positive_docid in collection

        if unjudged_docid in qrels[topicid]:
            stats['docid_collision'] += 1
            continue

        stats['kept'] += 1

        #if collection_length[positive_docid] > max_doc_token_length and collection_length[unjudged_docid] > max_doc_token_length:
        #    stats['both_to_long'] += 1
        #    continue
        #if collection_length[positive_docid] > max_doc_token_length:
        #    stats['pos_to_long'] += 1
        #    continue
        #if collection_length[unjudged_docid] > max_doc_token_length:
        #    stats['unjuged_to_long'] += 1
        #    continue

        triples.append((topicid,positive_docid,unjudged_docid))

# important: shuffle the train data
random.shuffle(triples)


with open(args.out_file,"w",encoding="utf8") as out_file_text ,\
     open(args.out_file_ids,"w",encoding="utf8") as out_file_ids:

    for i,(topicid, positive_docid, unjudged_docid) in tqdm(enumerate(triples)):
        if i == max_triples:
            break
        if collection[positive_docid].strip() != "" and collection[unjudged_docid].strip() != "":
            out_file_ids.write(str(topicid)+"\t"+positive_docid+"\t"+unjudged_docid+"\n")
            out_file_text.write(queries[topicid]+"\t"+collection[positive_docid]+"\t"+collection[unjudged_docid]+"\n")

for key, val in stats.items():
    print(f"{key}\t{val}")
[((6, 0, 6, 15), 'random.seed', 'random.seed', ({(6, 12, 6, 14): '(42)'}, {}), '(42)', False, 'import random\n'), ((21, 9, 21, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((55, 12, 55, 32), 'matchmaker.dataloaders.bling_fire_tokenizer.BlingFireTokenizer', 'BlingFireTokenizer', ({}, {}), '()', False, 'from matchmaker.dataloaders.bling_fire_tokenizer import BlingFireTokenizer\n'), ((81, 8, 81, 24), 'collections.defaultdict', 'defaultdict', ({(81, 20, 81, 23): 'int'}, {}), '(int)', False, 'from collections import defaultdict\n'), ((128, 0, 128, 23), 'random.shuffle', 'random.shuffle', ({(128, 15, 128, 22): 'triples'}, {}), '(triples)', False, 'import random\n'), ((12, 16, 12, 27), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((57, 16, 57, 37), 'tqdm.tqdm', 'tqdm', ({(57, 21, 57, 36): 'collection_file'}, {}), '(collection_file)', False, 'from tqdm import tqdm\n'), ((66, 16, 66, 32), 'tqdm.tqdm', 'tqdm', ({(66, 21, 66, 31): 'query_file'}, {}), '(query_file)', False, 'from tqdm import tqdm\n'), ((85, 16, 85, 36), 'tqdm.tqdm', 'tqdm', ({(85, 21, 85, 35): 'candidate_file'}, {}), '(candidate_file)', False, 'from tqdm import tqdm\n'), ((106, 25, 106, 54), 'random.choice', 'random.choice', ({(106, 39, 106, 53): 'qrels[topicid]'}, {}), '(qrels[topicid])', False, 'import random\n')]