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

dataset_path = '/data/cvprw/AIC23/dataset/origin/validation/'
save_path = '/data/cvprw/AIC23/dataset/val_result/detect_result/'

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

site_list = [file for file in os.listdir(dataset_path)]
temp = []
for site in site_list:
    d = os.path.join(dataset_path, site)
    if os.path.isdir(d) and not site.startswith('.'):
        temp.append(site)
site_list = natsort.natsorted(temp)


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

    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))
        h,w,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%60 == 0:
                # make folder
                if not os.path.exists(os.path.join(save_path,site_1,channel)):
                    os.makedirs(os.path.join(save_path,site_1,channel,"dets"))
                    os.makedirs(os.path.join(save_path,site_1,channel,"dets_debug"))
                    os.makedirs(os.path.join(save_path,site_1,channel,"labels"))

                cv2.imwrite(os.path.join(save_path,site_1,channel,"dets_debug",str(frame_count).zfill(5)+".png"),cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

                # based on label
                find_list = find_rows(list_file, frame_count)
                
                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])

                    # Crop
                    frame_arr = np.array(frame)
                    cropped_arr = frame_arr[y:y+h,x:x+w]

                    # Convert array to image
                    object_img = Image.fromarray(cropped_arr)

                    save_image_name = frame_n.zfill(5)+"_"+track_id.zfill(3)+".png"

                    object_img.save(os.path.join(save_path,site_1,channel,"dets",save_image_name))

            frame_count+=1
            ret, frame = cap.read()