import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
from IPython.display import clear_output
import os
from PIL import Image
import pandas as pd
import time
import natsort
import json

# site 별로 annotation file 만들자

dataset_path = '/data/cvprw/AIC23/dataset/origin/train/'
# save_path = '/data/cvprw/AIC23/dataset/coco_AIC23/'

site_list = [file for file in os.listdir(dataset_path)]
site_list = natsort.natsorted(site_list)

# for site in site_list:

# categoires part of coco
category_list = ['person','none']
category_dic_list=[]

def find_rows(list_file, x):
    find_bit = 0
    find_list = []
    for row in list_file:
        if int(row.split(',')[0]) == x:
            find_bit = 1
            find_list.append(row)
        else:
            if find_bit == 1:
                break
    
    return find_list

for idx, label in enumerate(category_list):
    category_dic={}
    category_dic["id"]=idx+1
    category_dic["name"]=label
    category_dic["supercategory"]='AIC23'
    category_dic_list.append(category_dic)
    
for site_1 in site_list:
    print(site_1)
    
    save_path = os.path.join("/data/cvprw/AIC23/dataset/coco_AIC23/train",site_1)
    save_name = "train.json"
    
    if not os.path.exists(save_path): 
        os.makedirs(save_path)
        
    coco_format={}
    coco_format['images']=[]
    coco_format['annotations']=[]

    temp = []


    channel_list = [file for file in os.listdir(os.path.join(dataset_path,site_1))]
    for channel in channel_list:
        d = os.path.join(dataset_path, site_1, channel)
        if os.path.isdir(d):
    #         print(channel)
            temp.append(channel)

    # var for coco
    total_n=1
    total_img_id=1

    for channel in temp:
        video_path = os.path.join(os.path.join(dataset_path, site_1, channel,'video.mp4'))
        label_path = os.path.join(os.path.join(dataset_path, site_1, channel,'label.txt'))

        ## Video
        cap = cv2.VideoCapture(video_path)
        ret, frame = cap.read()   

        video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_fps = int(cap.get(cv2.CAP_PROP_FPS))
        img_height,img_width,c = frame.shape

        # print(video_length)
        # print(video_fps)
        # print([h,w,c])

        # plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        # plt.show()

        ## Label
        with open(label_path, 'r') as f:
            list_file = f.readlines()
        list_file = [line.rstrip('\n') for line in list_file] 

        frame_count = 0
        while frame is not None:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            if frame_count != 0 and frame_count%60 == 0:
            # S002/c008/dets_debug/00000.png

                image_dic={}
                image_dic['id']=total_img_id
                image_dic['width']=int(img_width)
                image_dic['height']=int(img_height)
                image_dic['file_name']=os.path.join(site_1, channel, 'dets_debug',str(frame_count).zfill(5)+'.png')
                coco_format['images'].append(image_dic)

                find_list = find_rows(list_file, frame_count)

                # based on label
                for item in find_list:
                    row = item.split(',')

                    frame_n = row[0]
                    track_id = row[1]
                    x = int(row[2])
                    y = int(row[3])
                    w = int(row[4])
                    h = int(row[5])

                    annoation_dic={} 
                    annoation_dic['id']=total_n
                    annoation_dic['image_id']=total_img_id
                    annoation_dic['category_id']=1
                    annoation_dic['segmentation']=[[]]

                    # x,y,w,h
                    annoation_dic['bbox']=[x,y,w,h]
                    annoation_dic['area']=w*h
                    annoation_dic['iscrowd']=0
                    coco_format['annotations'].append(annoation_dic)

                    total_n += 1

                total_img_id+=1
            frame_count+=1
            ret, frame = cap.read()
            
    # categories
    coco_format['categories']=category_dic_list

    #save json
    with open(os.path.join(save_path,save_name), 'w', encoding='utf-8') as make_file:
        json.dump(coco_format, make_file, indent="\t")