import json
import random

import pandas as pd
from config import class_dict_A
import os

class_name_to_idx = dict((v, k) for k, v in class_dict_A.items())

anno_file = open('gt.json', 'r')
gt_file = json.load(anno_file)
anno_file.close()

split_file = open('split_train.txt', 'r')
vid_list = []
for line in split_file:
    vid_list.append(line.strip())
split_file.close()

class_lst = []
cls_idx_lst = []
start_lst = []
stop_lst = []
video_id_lst = []
point_lst = []
for vid_name in vid_list:
    vid_gt = gt_file['database'][vid_name]
    vid_fps = vid_gt['fps']
    t_factor = 25 / (vid_fps * 16)

    vid_annos = vid_gt['annotations']
    class_single_lst = []
    cls_idx_single_lst = []
    start_single_lst = []
    stop_single_lst = []
    video_id_single_lst = []
    point_single_lst = []
    for vid_anno in vid_annos:
        class_seg = vid_anno['label']
        class_idx = class_name_to_idx[class_seg]
        segment = vid_anno['segment']
        start_frame = round(segment[0] * vid_fps)
        stop_frame = round(segment[1] * vid_fps)
        random_point = random.randint(start_frame, int(start_frame + 0.6 * (stop_frame - start_frame)))

        class_single_lst.append(class_seg)
        cls_idx_single_lst.append(class_idx)
        start_single_lst.append(start_frame)
        stop_single_lst.append(stop_frame)
        video_id_single_lst.append(vid_name)
        point_single_lst.append(random_point)

    class_lst.extend(class_single_lst)
    cls_idx_lst.extend(cls_idx_single_lst)
    start_lst.extend(start_single_lst)
    stop_lst.extend(stop_single_lst)
    video_id_lst.extend(video_id_single_lst)
    point_lst.extend(point_single_lst)

# 字典中的key值即为csv中的列名
frame = pd.DataFrame({'class': class_lst, 'class_index': cls_idx_lst, 'start_frame': start_lst, 'stop_frame': stop_lst, 'video_id': video_id_lst, 'point': point_lst})

# 将DataFrame存储为csv，index表识时候显示行名，default=True
frame.to_csv("ex1.csv", index=False, sep=',')
