import argparse
from trec_car import read_data
from tqdm import tqdm
import pickle
import os
import json
import copy
from utils.util import NUM_FOLD


def parse_sim_file(filename):
    """
    Reads the deduplicated documents file and stores the
    duplicate passage ids into a dictionary
    """

    sim_dict = {}
    lines = open(filename).readlines()
    for line in lines:
        data = line.strip().split(':')
        if len(data[1]) > 0:
            sim_docs = data[-1].split(',')
            for docs in sim_docs:
                sim_dict[docs] = 1

    return sim_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--car_cbor", type=str)
    parser.add_argument("--msmarco_collection", type=str)
    parser.add_argument("--duplicate_file", type=str)
    parser.add_argument("--cast_dir", type=str)

    parser.add_argument("--out_data_dir", type=str)
    parser.add_argument("--out_collection_dir", type=str)
    args = parser.parse_args()

    # INPUT
    sim_file = args.duplicate_file
    cast_topics_raw_file = os.path.join(args.cast_dir,
                                        "evaluation_topics_v1.0.json")
    cast_topics_manual_file = os.path.join(
        args.cast_dir, "evaluation_topics_annotated_resolved_v1.0.tsv")
    cast_qrels_file = os.path.join(args.cast_dir, "2019qrels.txt")

    # OUTPUT
    out_topics_file = os.path.join(args.out_data_dir, "eval_topics.jsonl")
    out_raw_queries_file = os.path.join(args.out_data_dir, "queries.raw.tsv")
    out_manual_queries_file = os.path.join(args.out_data_dir,
                                           "queries.manual.tsv")
    out_qrels_file = os.path.join(args.out_data_dir, "qrels.tsv")
    car_id_to_idx_file = os.path.join(args.out_collection_dir,
                                      "car_id_to_idx.pickle")
    car_idx_to_id_file = os.path.join(args.out_collection_dir,
                                      "car_idx_to_id.pickle")
    out_collection_file = os.path.join(args.out_collection_dir,
                                       "collection.tsv")


    # 1. Combine TREC-CAR & MS MARCO, remove duplicate passages, assign new ids
    car_id_to_idx = {}
    car_idx_to_id = []
    if os.path.exists(out_collection_file) and os.path.exists(
            car_id_to_idx_file) and os.path.exists(car_idx_to_id_file):
        print("Preprocessed collection found. Loading car_id_to_idx...")
        with open(car_id_to_idx_file, "rb") as f:
            car_id_to_idx = pickle.load(f)
    else:
        sim_dict = parse_sim_file(sim_file)
        car_base_id = 10000000
        i = 0
        with open(out_collection_file, "w", encoding="utf-8") as f: #FIX change 'a' to 'w' in normal run
            print("Processing TREC-CAR...")
            for para in tqdm(
                    read_data.iter_paragraphs(open(args.car_cbor, 'rb'))):
                car_id = "CAR_" + para.para_id
                text = para.get_text()
                text = text.replace("\t", " ").replace("\n",
                                                       " ").replace("\r", " ")
                idx = car_base_id + i
                car_id_to_idx[
                    car_id] = idx  # e.g. CAR_76a4a716d4b1b01995c6663ee16e94b4ca35fdd3 -> 10000044
                car_idx_to_id.append(car_id)
                f.write("{}\t{}\n".format(idx, text))
                i += 1
            print("Processing MS MARCO...")
            removed = 0
            with open(args.msmarco_collection, "r") as m:
                for line in tqdm(m):
                    marco_id, text = line.strip().split("\t")
                    if ("MARCO_" + marco_id) in sim_dict:
                        removed += 1
                        continue
                    f.write("{}\t{}\n".format(marco_id, text))
            print("Removed " + str(removed) + " passages")
        print("Dumping id mappings to {} and {}...".format(car_id_to_idx_file, car_idx_to_id_file))
        with open(car_id_to_idx_file, "wb") as f:
            pickle.dump(car_id_to_idx, f)
        with open(car_idx_to_id_file, "wb") as f:
            pickle.dump(car_idx_to_id, f)

    # 2. Process queries
    print("Processing CAsT utterances...")
    with open(cast_topics_raw_file, "r") as fin:
        raw_data = json.load(fin)

    with open(cast_topics_manual_file, "r") as fin:
        annonated_lines = fin.readlines()

    out_raw_queries = open(out_raw_queries_file, "w")
    out_manual_queries = open(out_manual_queries_file, "w")

    all_annonated = {}
    for line in annonated_lines:
        splitted = line.split('\t')
        out_manual_queries.write(line)
        topic_query = splitted[0]
        query = splitted[1].strip()
        topic_id = topic_query.split('_')[0]
        query_id = topic_query.split('_')[1]
        if topic_id not in all_annonated:
            all_annonated[topic_id] = {}
        all_annonated[topic_id][query_id] = query
    out_manual_queries.close()

    topic_number_dict = {}
    data = []
    for group in raw_data:
        topic_number, description, turn, title = str(
            group['number']), group.get('description',
                                        ''), group['turn'], group.get(
                                            'title', '')
        queries = []
        for query in turn:
            query_number, raw_utterance = str(
                query['number']), query['raw_utterance']
            queries.append(raw_utterance)
            record = {}
            record['topic_number'] = topic_number
            record['query_number'] = query_number
            record['description'] = description
            record['title'] = title
            record['input'] = copy.deepcopy(queries)
            record['target'] = all_annonated[topic_number][query_number]
            out_raw_queries.write("{}_{}\t{}\n".format(topic_number,
                                                       query_number,
                                                       raw_utterance))
            if not topic_number in topic_number_dict:
                topic_number_dict[topic_number] = len(topic_number_dict)
            data.append(record)
    out_raw_queries.close()

    with open(out_topics_file, 'w') as fout:
        for item in data:
            json_str = json.dumps(item)
            fout.write(json_str + '\n')

    # Split eval data into K-fold
    topic_per_fold = len(topic_number_dict) // NUM_FOLD
    for i in range(NUM_FOLD):
        with open(out_topics_file + "." + str(i), 'w') as fout:
            for item in data:
                idx = topic_number_dict[item['topic_number']]
                if idx // topic_per_fold == i:
                    json_str = json.dumps(item)
                    fout.write(json_str + '\n')

    # 3. Process and convert qrels
    print("Processing qrels...")
    with open(cast_qrels_file, "r") as oq, open(out_qrels_file, "w") as nq:
        for line in oq:
            qid, _, pid, rel = line.strip().split()
            if pid.startswith("CAR_"):
                assert car_id_to_idx[pid] != -1
                pid = car_id_to_idx[pid]
            elif pid.startswith("MARCO_"):
                pid = int(pid[6:])
            else:
                continue
            nq.write(qid + "\t0\t" + str(pid) + "\t" + rel + "\n")

    print("End")

[((30, 13, 30, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((42, 27, 43, 70), 'os.path.join', 'os.path.join', ({(42, 40, 42, 53): 'args.cast_dir', (43, 40, 43, 69): '"""evaluation_topics_v1.0.json"""'}, {}), "(args.cast_dir, 'evaluation_topics_v1.0.json')", False, 'import os\n'), ((44, 30, 45, 71), 'os.path.join', 'os.path.join', ({(45, 8, 45, 21): 'args.cast_dir', (45, 23, 45, 70): '"""evaluation_topics_annotated_resolved_v1.0.tsv"""'}, {}), "(args.cast_dir, 'evaluation_topics_annotated_resolved_v1.0.tsv')", False, 'import os\n'), ((46, 22, 46, 66), 'os.path.join', 'os.path.join', ({(46, 35, 46, 48): 'args.cast_dir', (46, 50, 46, 65): '"""2019qrels.txt"""'}, {}), "(args.cast_dir, '2019qrels.txt')", False, 'import os\n'), ((49, 22, 49, 74), 'os.path.join', 'os.path.join', ({(49, 35, 49, 52): 'args.out_data_dir', (49, 54, 49, 73): '"""eval_topics.jsonl"""'}, {}), "(args.out_data_dir, 'eval_topics.jsonl')", False, 'import os\n'), ((50, 27, 50, 77), 'os.path.join', 'os.path.join', ({(50, 40, 50, 57): 'args.out_data_dir', (50, 59, 50, 76): '"""queries.raw.tsv"""'}, {}), "(args.out_data_dir, 'queries.raw.tsv')", False, 'import os\n'), ((51, 30, 52, 64), 'os.path.join', 'os.path.join', ({(51, 43, 51, 60): 'args.out_data_dir', (52, 43, 52, 63): '"""queries.manual.tsv"""'}, {}), "(args.out_data_dir, 'queries.manual.tsv')", False, 'import os\n'), ((53, 21, 53, 65), 'os.path.join', 'os.path.join', ({(53, 34, 53, 51): 'args.out_data_dir', (53, 53, 53, 64): '"""qrels.tsv"""'}, {}), "(args.out_data_dir, 'qrels.tsv')", False, 'import os\n'), ((54, 25, 55, 61), 'os.path.join', 'os.path.join', ({(54, 38, 54, 61): 'args.out_collection_dir', (55, 38, 55, 60): '"""car_id_to_idx.pickle"""'}, {}), "(args.out_collection_dir, 'car_id_to_idx.pickle')", False, 'import os\n'), ((56, 25, 57, 61), 'os.path.join', 'os.path.join', ({(56, 38, 56, 61): 'args.out_collection_dir', (57, 38, 57, 60): '"""car_idx_to_id.pickle"""'}, {}), "(args.out_collection_dir, 'car_idx_to_id.pickle')", False, 'import os\n'), ((58, 26, 59, 56), 'os.path.join', 'os.path.join', ({(58, 39, 58, 62): 'args.out_collection_dir', (59, 39, 59, 55): '"""collection.tsv"""'}, {}), "(args.out_collection_dir, 'collection.tsv')", False, 'import os\n'), ((65, 7, 65, 42), 'os.path.exists', 'os.path.exists', ({(65, 22, 65, 41): 'out_collection_file'}, {}), '(out_collection_file)', False, 'import os\n'), ((65, 47, 66, 31), 'os.path.exists', 'os.path.exists', ({(66, 12, 66, 30): 'car_id_to_idx_file'}, {}), '(car_id_to_idx_file)', False, 'import os\n'), ((66, 36, 66, 70), 'os.path.exists', 'os.path.exists', ({(66, 51, 66, 69): 'car_idx_to_id_file'}, {}), '(car_idx_to_id_file)', False, 'import os\n'), ((107, 19, 107, 33), 'json.load', 'json.load', ({(107, 29, 107, 32): 'fin'}, {}), '(fin)', False, 'import json\n'), ((69, 28, 69, 42), 'pickle.load', 'pickle.load', ({(69, 40, 69, 41): 'f'}, {}), '(f)', False, 'import pickle\n'), ((100, 12, 100, 41), 'pickle.dump', 'pickle.dump', ({(100, 24, 100, 37): 'car_id_to_idx', (100, 39, 100, 40): 'f'}, {}), '(car_id_to_idx, f)', False, 'import pickle\n'), ((102, 12, 102, 41), 'pickle.dump', 'pickle.dump', ({(102, 24, 102, 37): 'car_idx_to_id', (102, 39, 102, 40): 'f'}, {}), '(car_idx_to_id, f)', False, 'import pickle\n'), ((145, 30, 145, 52), 'copy.deepcopy', 'copy.deepcopy', ({(145, 44, 145, 51): 'queries'}, {}), '(queries)', False, 'import copy\n'), ((157, 23, 157, 39), 'json.dumps', 'json.dumps', ({(157, 34, 157, 38): 'item'}, {}), '(item)', False, 'import json\n'), ((91, 28, 91, 35), 'tqdm.tqdm', 'tqdm', ({(91, 33, 91, 34): 'm'}, {}), '(m)', False, 'from tqdm import tqdm\n'), ((167, 31, 167, 47), 'json.dumps', 'json.dumps', ({(167, 42, 167, 46): 'item'}, {}), '(item)', False, 'import json\n')]