在# -*-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 in ['dvSave-2022_01_26_20_01_46', 'dvSave-2022_03_21_11_09_50',
                        'dvSave-2022_03_21_16_11_40', 'dvSave-2022_03_21_11_12_27',
                        'dvSave-2022_03_21_11_19_37', 'dvSave-2022_02_25_13_21_41',
                        'dvSave-2022_02_25_13_31_18', 'dvSave-2022_03_21_09_05_49']:
            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.")
，执行报：xtracting frame timestamps…
Frame timestamp start of frame: 1630493187262566 end of frame: 1630493187297988
Frame timestamp start of exposure: 1630493187271725 end of exposure: 1630493187288806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187302566 end of frame: 1630493187337988
Frame timestamp start of exposure: 1630493187311725 end of exposure: 1630493187328806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187342566 end of frame: 1630493187377988
Frame timestamp start of exposure: 1630493187351725 end of exposure: 1630493187368806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187382566 end of frame: 1630493187417988
Frame timestamp start of exposure: 1630493187391725 end of exposure: 1630493187408806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187422566 end of frame: 1630493187457988
Frame timestamp start of exposure: 1630493187431725 end of exposure: 1630493187448806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187462566 end of frame: 1630493187497988
Frame timestamp start of exposure: 1630493187471725 end of exposure: 1630493187488806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187502566 end of frame: 1630493187537988
Frame timestamp start of exposure: 1630493187511725 end of exposure: 1630493187528806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187542566 end of frame: 1630493187577988
Frame timestamp start of exposure: 1630493187551725 end of exposure: 1630493187568806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187582566 end of frame: 1630493187617988
Frame timestamp start of exposure: 1630493187591725 end of exposure: 1630493187608806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187622566 end of frame: 1630493187657988
Frame timestamp start of exposure: 1630493187631725 end of exposure: 1630493187648806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187662566 end of frame: 1630493187697988
Frame timestamp start of exposure: 1630493187671725 end of exposure: 1630493187688806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187702566 end of frame: 1630493187737989
Frame timestamp start of exposure: 1630493187711725 end of exposure: 1630493187728806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187742566 end of frame: 1630493187777989
Frame timestamp start of exposure: 1630493187751725 end of exposure: 1630493187768806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187782566 end of frame: 1630493187817989
Frame timestamp start of exposure: 1630493187791725 end of exposure: 1630493187808806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187822567 end of frame: 1630493187857989
Frame timestamp start of exposure: 1630493187831725 end of exposure: 1630493187848806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187862567 end of frame: 1630493187897989
Frame timestamp start of exposure: 1630493187871725 end of exposure: 1630493187888806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187902567 end of frame: 1630493187937989
Frame timestamp start of exposure: 1630493187911725 end of exposure: 1630493187928806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187942567 end of frame: 1630493187977989
Frame timestamp start of exposure: 1630493187951725 end of exposure: 1630493187968806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493187982567 end of frame: 1630493188017989
Frame timestamp start of exposure: 1630493187991725 end of exposure: 1630493188008806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188022567 end of frame: 1630493188057989
Frame timestamp start of exposure: 1630493188031725 end of exposure: 1630493188048806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188062567 end of frame: 1630493188097989
Frame timestamp start of exposure: 1630493188071725 end of exposure: 1630493188088806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188102567 end of frame: 1630493188137989
Frame timestamp start of exposure: 1630493188111725 end of exposure: 1630493188128806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188142567 end of frame: 1630493188177989
Frame timestamp start of exposure: 1630493188151725 end of exposure: 1630493188168806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188182567 end of frame: 1630493188217989
Frame timestamp start of exposure: 1630493188191725 end of exposure: 1630493188208806
Extracting frame timestamps…
Frame timestamp start of frame: 1630493188222567 end of frame: 1630493188257989
Frame timestamp start of exposure: