import json
import argparse




def parse_args():
    parser = argparse.ArgumentParser("D2 model converter")

    parser.add_argument("--results", default="", type=str, help="Path to the GLEE output dir ")
    parser.add_argument("--refer", default="", type=str, help="Path to the BURST annotation val dir")

    return parser.parse_args()


def main():
    args = parse_args()

    ori_anno = json.load(open(args.results,'rb'))

    reference_anno = json.load(open(args.refer,'rb'))

    
    num_tracks = 0
    num_miss_video = 0


    id_mapping = {}

    for i, cate_info in enumerate(reference_anno['categories']):
        new_id = i
        old_id = cate_info['id']
        id_mapping.update({new_id:old_id})  


    ref_sequences_dict = {}
    for ref in reference_anno['sequences']:
        ref_sequences_dict[ref['id']] = ref


    # ids = [v['category_id'] for v in ori_anno]

    sequences_dict = {}

    for seg in ori_anno:
        vid = seg['video_id']
        if vid not in sequences_dict.keys():
            # import pdb;pdb.set_trace()
            sequences_dict[vid] = {
                'id': vid,
                'width': ref_sequences_dict[vid]['width'],
                'height': ref_sequences_dict[vid]['height'],
                'seq_name': ref_sequences_dict[vid]['seq_name'], 
                'dataset': ref_sequences_dict[vid]['dataset'],
                'annotated_image_paths': ref_sequences_dict[vid]['annotated_image_paths'],
                'fps': ref_sequences_dict[vid]['fps'],
                'segmentations': [{} for i in range(len(seg['segmentations']))],
                'track_category_ids': {},
            }
        track_id = str(len(sequences_dict[vid]['track_category_ids']) + 1)

        for frame, rles in enumerate(seg['segmentations']): 
            sequences_dict[vid]['segmentations'][frame][track_id] = {'rle': rles['counts'], 'score':seg['score']}
        # import pdb;pdb.set_trace()
        sequences_dict[vid]['track_category_ids'][track_id] = id_mapping[seg['category_id']]


    results = {'sequences':[]} 
    for k,v in sequences_dict.items():
        results['sequences'].append(v)

    
    with open('converted_tao_results.json', 'w') as f:
        json.dump(results, f)



if __name__ == "__main__":
    main()