import json
from collections import defaultdict

splits = ['train', 'val_seen', 'val_unseen']
# splits = ['train']
# for split in splits:
#     print(split)
#     new_data = []
#     with open('../RXR/rxr-data/rxr_%s_guide.jsonl' % split, 'r') as f:
#         for line in f:
#             new_data.append(json.loads(line))
#     f.close()
#
#     combined_data = defaultdict(list)
#     for i, data in enumerate(new_data):
#         combined_data[data['path_id']].append(data)
#
#     final_data = []
#     for i, data in enumerate(new_data):
#         print(i)
#         data['instructions'] = []
#         data['languages'] = []
#         data['instruction_ids'] = []
#         for ins in combined_data[data['path_id']]:
#             if ins['instruction_id'] == data['instruction_id']:
#                 assert data['language'] == ins['language']
#                 continue
#             data['instructions'].append(ins['instruction'])
#             data['languages'].append(ins['language'])
#             data['instruction_ids'].append(ins['instruction_id'])
#         final_data.append(data)
#
#     with open('../RXR/rxr-data/rxr_%s_guide_multi.jsonl' % split, 'w') as f:
#         json.dump(final_data, f)
#     f.close()

# for split in splits:
#     print(split)
#     new_data = []
#     with open('../RXR/rxr-data/rxr_%s_guide_multi.jsonl' % split, 'r') as f:
#         new_data = json.load(f)
#     f.close()
#
#     follower = []
#     with open('../RXR/rxr-data/rxr_%s_follower.jsonl' % split, 'r') as f:
#         for line in f:
#             follower.append(json.loads(line))
#     f.close()
#
#     ins2path = dict()
#     for i, data in enumerate(follower):
#         ins2path[data['instruction_id']] = data['path']
#
#     for i, data in enumerate(new_data):
#         new_data[i]['follower_path'] = ins2path[data['instruction_id']]
#
#     with open('../RXR/rxr-data/rxr_%s_f+g_multi.jsonl' % split, 'w') as f:
#         json.dump(new_data, f)
#     f.close()


# new_data = []
# with open('../RXR/rxr-data/rxr_train_guide_multi.jsonl', 'r') as f:
#     new_data = json.load(f)
# f.close()
#
# for data in new_data:
#     if len(data['instructions']) != 8:
#         print(data)

# with open('../RXR/rxr-data/rxr_train_guide_multi.jsonl', 'r') as f:
#     new_data = json.load(f)
# f.close()

# for data in new_data:
#     if data['instruction_id'] == 55603:
#         print(data)
#         break

# for split in splits:
#     print(split)
#     with open('../RXR/rxr-data/rxr_%s_guide_multi.jsonl' % split, 'r') as f:
#         new_data = json.load(f)
#     f.close()
#     scan = []
#     for i, data in enumerate(new_data):
#         scan.append(data['scan'])
#     print(set(scan))

with open('objects_lvis_36.json', 'r') as f:
    objects = json.load(f)
f.close()

# from detectron2 import model_zoo
# from detectron2.engine import DefaultPredictor
# from detectron2.config import get_cfg
# from detectron2.utils.visualizer import Visualizer
# from detectron2.data import MetadataCatalog, DatasetCatalog
#
# from collections import defaultdict


# dataset = 'lvis'
#
# cfg = get_cfg()
#
# if dataset == 'coco':
#     cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
#     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
#     cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
# elif dataset == 'lvis':
#     cfg.merge_from_file(model_zoo.get_config_file("LVISv0.5-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml"))
#     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
#     cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("LVISv0.5-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml")
# else:
#     print("Error")
#     exit()
#
# attribute = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes
# print(len(attribute))

# attribute = ['painting', 'cabinet', 'chair', 'bed', 'sofa', 'sculpture', 'refrigerator', 'sink', 'toilet', 'fireplace', 'table', 'ladder']

attribute = ['drawer', 'faucet', 'cabinet', 'hinge', 'cushion', 'sofa', 'chair', 'pillow', 'armchair', 'lamp', 'vase', 'knob', 'curtain', 'statue_(sculpture)', 'doorknob', 'vent', 'lightbulb', 'flowerpot', 'book', 'pipe', 'painting', 'wall_socket', 'bed', 'mirror', 'television_set', 'flower_arrangement', 'chandelier']

# 27 attributes

attributes = dict()

for scan, v in objects.items():
    attributes[scan] = dict()
    for view, i in v.items():
        attributes[scan][view] = dict()
        for img, name in i.items():
            result = []
            for a in attribute:
                if a in name:
                    result.append(1)
                else:
                    result.append(0)
            attributes[scan][view][img] = result

with open('objects_27.json', 'w') as f:
    json.dump(attributes, f)
f.close()

# for split in splits:
#     print(split)
#     with open('../RXR/rxr-data/rxr_%s_guide_multi.jsonl' % split, 'r') as f:
#         new_data = json.load(f)
#     f.close()
#
#     ins2sv = dict()
#     for i, data in enumerate(new_data):
#         instr_id = data['instruction_id']
#         ins2sv[instr_id] = [data['scan'], data]




