# -*-coding:utf-8-*-
import torch
from utils.loading_utils import load_model, get_device
import numpy as np
import argparse
import pandas as pd
import os
import cv2
import pdb
from utils.event_readers import FixedSizeEventReader, FixedDurationEventReader
from utils.inference_utils import events_to_voxel_grid, events_to_voxel_grid_pytorch
from utils.timers import Timer
import time
from image_reconstructor import ImageReconstructor
from options.inference_options import set_inference_options
from PIL import Image
from dv import AedatFile



def extract_davis(aedat_file_path, dvs_img_interval):
    frame_no = 0
    frame_all = []
    frame_exposure_time = []
    frame_interval_time = []
    use_mode = 'frame_interval_time'
    
    with AedatFile(aedat_file_path) as f:
        print(f.names)  # List all streams
        
        for frame in f['frames']:
            print("Extracting frame timestamps…")
            frame_all.append(frame.image)
            frame_exposure_time.append([frame.timestamp_start_of_exposure,
                                        frame.timestamp_end_of_exposure])
            frame_interval_time.append([frame.timestamp_start_of_frame, frame.timestamp_end_of_frame])

            print("Frame timestamp start of frame:", frame.timestamp_start_of_frame, 
                  "end of frame:", frame.timestamp_end_of_frame)
            print("Frame timestamp start of exposure:", frame.timestamp_start_of_exposure, 
                  "end of exposure:", frame.timestamp_end_of_exposure)
        
        if use_mode == 'frame_exposure_time':
            frame_timestamp = frame_exposure_time
        elif use_mode == 'frame_interval_time':
            frame_timestamp = frame_interval_time

        frame_num = len(frame_timestamp)
        events = np.hstack([packet for packet in f['events'].numpy()])
        
        # Print out the min and max timestamps of the events
        print("Events timestamp range:", events['timestamp'].min(), "to", events['timestamp'].max())
        print("Frame timestamp range:", frame_timestamp[0], "to", frame_timestamp[-1])
        
        return events['timestamp'], events['x'], events['y'], events['polarity'], frame_num, frame_timestamp


    #######################################################################################################


####                                Main Function
#######################################################################################################
if __name__ == "__main__":
    device = torch.device("cuda:4")
    # data_path = r"/home/ioe/xxxxx/EventTracking/COESOT/train"
    data_path = r"/data2/local_userdata/heyuting/datasets/COESOT/test"
    save_path = r"/data2/local_userdata/heyuting/datasets/COESOT/test"
    video_files = os.listdir(data_path)
    dvs_img_interval = 1
    for videoID in range(len(video_files)):
        foldName = video_files[videoID]
        #print("==>> foldName: ", foldName)
        if foldName not in ['dvSave-2022_08_21_13_21_14']:
            #print("==>> pass it: ", foldName)
            continue
        aedat4_file = foldName + '.aedat4'
        aedat_file_path = os.path.join(data_path, foldName, aedat4_file)
        # aedat_file_path = os.path.join(data_path, foldName, foldName+'.aedat4')
        save_img_path = os.path.join(save_path, foldName, foldName+'_recon')
        if not os.path.exists(save_img_path):
            os.makedirs(save_img_path)


        timestamps, x, y, polarities, frame_num, frame_timestamp = extract_davis(aedat_file_path, 1)  ## (4391964,)

        # all_events = [timestamps, x, y, polarities]

        all_events = np.zeros((timestamps.shape[0], 4))

        all_events[:, 0] = timestamps
        all_events[:, 1] = x
        all_events[:, 2] = y
        all_events[:, 3] = polarities
        # pdb.set_trace()

        parser = argparse.ArgumentParser(description='Evaluating a trained network')
        parser.add_argument('-c', '--path_to_model', required=True, type=str, help='path to model weights')
        parser.add_argument('-i', '--input_file', required=False, type=str)
        parser.add_argument('--fixed_duration', dest='fixed_duration', action='store_true')
        parser.set_defaults(fixed_duration=False)
        parser.add_argument('-N', '--window_size', default=500000, type=int,
                            help="Size of each event window, in number of events. Ignored if --fixed_duration=True")
        parser.add_argument('-T', '--window_duration', default=33.33, type=float,
                            help="Duration of each event window, in milliseconds. Ignored if --fixed_duration=False")
        parser.add_argument('--num_events_per_pixel', default=0.35, type=float,
                            help='in case N (window size) is not specified, it will be automatically computed as N = width * height * num_events_per_pixel')
        parser.add_argument('--skipevents', default=0, type=int)
        parser.add_argument('--suboffset', default=0, type=int)
        parser.add_argument('--compute_voxel_grid_on_cpu', dest='compute_voxel_grid_on_cpu',
                            action='store_true')
        parser.set_defaults(compute_voxel_grid_on_cpu=False)

        set_inference_options(parser)

        args = parser.parse_args()

        # Load model
        device = get_device(args.use_gpu)
        #model = load_model(args.path_to_model, device)
        model = load_model(args.path_to_model)
        model = model.to(device)
        model.eval()

        # pdb.set_trace()
        height, width = 260, 346
        reconstructor = ImageReconstructor(model, height, width, model.num_bins, args)
        start_s = 0
        """ Read chunks of events using Pandas """
        for frame_no in range(0, frame_num - 1):
            N = frame_timestamp[frame_no][1] - frame_timestamp[frame_no][0]
            initial_offset = args.skipevents
            sub_offset = args.suboffset
            start_index = initial_offset + sub_offset
            ##############################
            try:
                print(f"Processing frame {frame_no}, timestamp range: {frame_timestamp[frame_no]}")

                start_idx = np.searchsorted(all_events[:, 0], frame_timestamp[frame_no][0], side='left')
                end_idx = np.searchsorted(all_events[:, 0], frame_timestamp[frame_no][1], side='right')
                print("Frame", frame_no, "Start index:", start_idx, "End index:", end_idx)

                if start_idx < end_idx:
                    event_window = all_events[start_idx:end_idx]
                    if event_window.shape[0] == 0:
                        print("Warning: Event window is empty.")
                        continue
                    last_timestamp = event_window[-1, 0]

                    with Timer('Processing entire dataset'):
                        if args.compute_voxel_grid_on_cpu:
                            event_tensor = events_to_voxel_grid(event_window, num_bins=model.num_bins, width=width, height=height)
                            event_tensor = torch.from_numpy(event_tensor)
                        else:
                            event_tensor = events_to_voxel_grid_pytorch(event_window, num_bins=model.num_bins, width=width, height=height, device=device)

                    num_events_in_window = event_window.shape[0]
                    reconstructor.update_reconstruction(event_tensor, start_index + num_events_in_window, last_timestamp, frame_no, save_img_path)
                    start_index += num_events_in_window
                else:
                    print("Warning: Invalid index range. Start idx >= End idx.")
            except IndexError as e:
                print(f"IndexError encountered: {e}, skipping this frame.")
，报：Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059288886426 end of exposure: 1661059288910429
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059288928771 end of exposure: 1661059288952775
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059288971117 end of exposure: 1661059288995120
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289013462 end of exposure: 1661059289037465
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289055807 end of exposure: 1661059289079810
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289098152 end of exposure: 1661059289122155
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289140497 end of exposure: 1661059289164500
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289182842 end of exposure: 1661059289206845
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289225186 end of exposure: 1661059289249190
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289267531 end of exposure: 1661059289291535
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289309878 end of exposure: 1661059289333880
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289352223 end of exposure: 1661059289376226
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289394568 end of exposure: 1661059289418571
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289436913 end of exposure: 1661059289460916
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289479258 end of exposure: 1661059289503261
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289521604 end of exposure: 1661059289545606
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289563949 end of exposure: 1661059289587951
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289606294 end of exposure: 1661059289630296
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289648639 end of exposure: 1661059289672641
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289690984 end of exposure: 1661059289714986
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289733329 end of exposure: 1661059289757331
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289775674 end of exposure: 1661059289799676
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289818019 end of exposure: 1661059289842022
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289860364 end of exposure: 1661059289884367
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289902708 end of exposure: 1661059289926712
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289945053 end of exposure: 1661059289969057
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059289987399 end of exposure: 1661059290011402
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290029743 end of exposure: 1661059290053747
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290072088 end of exposure: 1661059290096092
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290114433 end of exposure: 1661059290138437
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290156779 end of exposure: 1661059290180782
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290199124 end of exposure: 1661059290223127
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290241470 end of exposure: 1661059290265472
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290283815 end of exposure: 1661059290307818
Extracting frame timestamps…
Frame timestamp start of frame: 0 end of frame: 0
Frame timestamp start of exposure: 1661059290326160 end of exposure: 1661059290350163