import torch

import os
import time
import json
import numpy as np
from collections import defaultdict
from speaker import Speaker
from mbert import mBERT

from utils import read_vocab,write_vocab,build_vocab,Tokenizer,padding_idx,timeSince, read_img_features
import utils
from env import R2RBatch
from agent import Seq2SeqAgent
from eval import Evaluation
import warnings
warnings.filterwarnings("ignore")


from tensorboardX import SummaryWriter
from transformers import BertTokenizer

import sys
import random
import math

import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
from utils import padding_idx, add_idx, Tokenizer
import model
from collections import defaultdict

from transformers import BertModel, BertConfig, AdamW, get_linear_schedule_with_warmup

import heapq

# splits = ['train', 'val_seen', 'val_unseen']
splits = ['val_unseen']
# splits = ['train']
counts = dict()
for split in splits:
    print(split)
    count = 0
    with open('../RXR/rxr-data/rxr_%s_guide_multi.jsonl' % split) as f:
        new_data = json.load(f)
    f.close()

    with open("../RXR/rxr-data/rxr_%s_guide_image_exp1lang_top5_screen.jsonl" % split) as f:
        results = json.load(f)
    f.close()

    ins_to_data = dict()
    for data in new_data:
        ins_to_data[data['instruction_id']] = data

    count = 0

    for i, data in enumerate(new_data):
        print(i)
        ins_id = data['instruction_id']
        result = results[str(ins_id)]
        top5_path = []
        for j in range(5):
            if result[j][1] in ins_to_data:
                top5_path.append(ins_to_data[result[j][1]]['path_id'])
        if data['path_id'] in top5_path:
            count += 1
        # sim_score = result[0][0]
        # print(result)
        # print(ins_id)
        # sim_score = result[0][0]
        # if result[0][1] not in ins_to_data or sim_score < 0.99:
        #     similar_data = data
        #     count += 1
        # else:
        #     similar_data = ins_to_data[result[0][1]]

        # if result[0][1] not in ins_to_data:
        #     similar_data = data
        # else:
        #     similar_data = ins_to_data[result[0][1]]
        # new_data[i]['pair_scan'] = similar_data['scan']
        # new_data[i]['pair_path'] = similar_data['path']
        # ins_id = data['instruction_id']
        # if similar_data['path_id'] == data['path_id']:
        #     count += 1

        # if data['language'].split("-")[0] == 'en':
        #     print("scan", data['scan'])
        #     print("path",data['path'])
        #     print("path_id", data['path_id'])
        #     print("ins", data['instruction'])
        #     # print("image", results[str(ins_id)])
        #     # result = results[str(ins_id)]
        #     similar_data1 = ins_to_data[int(result[0][1])]
        #     similar_data2 = ins_to_data[int(result[1][1])]
        #     print("s_scan1", similar_data1['scan'])
        #     print("s_path1", similar_data1['path'])
        #     print("s_ins1", similar_data1['instruction'])
        #     print("s_scan2", similar_data2['scan'])
        #     print("s_path2", similar_data2['path'])
        #     print("s_ins2", similar_data2['instruction'])
        #     count += 1
        #     if count > 10:
        #         exit()

    print(count / len(new_data))
    # print(split, count / len(new_data))
    # counts[split] = count / len(new_data)
    # exit()

    # with open('../RXR/rxr-data/rxr_%s_guide_multi_mbert.jsonl' % split, "w") as f:
    #     json.dump(new_data, f)
    # f.close()
# #
# print(counts)
# with open('../RXR/rxr-data/rxr_train_guide_multi.jsonl') as f:
#     new_data = json.load(f)
# f.close()
#
# count = 0
# for i, data in enumerate(new_data):
#     if len(data['path']) == 3:
#         print(data['path'])
#         print(data['scan'])
#         print(data['instruction'])
#         count += 1
#
#     if count > 3:
#         exit()
