import numpy as np
import cv2
import esim_py


# camera = cv2.VideoCapture(0)
camera = cv2.VideoCapture('video.mp4')  

POS_THRESHOLD = 0.5
NEG_THRESHOLD = 0.5
REF_PERIOD = 0.000

esim = esim_py.EventSimulator(POS_THRESHOLD, NEG_THRESHOLD, REF_PERIOD, 1e-4, True)  

# # generate events from list of images and timestamps
# events_list_of_images = esim.generateFromStampedImageSequence(
#     list_of_image_files,   # list of absolute paths to images
#     list_of_timestamps     # list of timestamps in ascending order
# )

fps = cv2.CAP_PROP_FPS
ts_s = 1 / fps
ts_ns = ts_s * 1e9 # convert s to ns

is_init = False
idx = 0
while True:
    _, frame_bgr = camera.read()
    frame_gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
    frame_log = np.log(frame_gray.astype("float32") / 255 + 1e-4)
    height, width = frame_log.shape[:2]

    current_ts_ns = idx * ts_ns

    if not is_init:
        esim.init(frame_log, current_ts_ns)
        is_init = True
        idx += 1
        
        continue
    
    events = esim.generateEventFromCVImage(frame_log, current_ts_ns)
    x, y, t, p = events.T
    t = t.astype(dtype=np.float32) * 1e-6 # convert ns to milliseconds
    
    last_timestamp = t[-1]

    event_frame = np.zeros((height, width, 3), dtype=np.float32)
    
    x = x.astype(dtype=np.int32)
    y = y.astype(dtype=np.int32)
    p = p.astype(dtype=np.int32)

    print(idx, events.shape)

    if last_timestamp <= 0:
        continue

    event_frame[y, x, 1 - p] = (last_timestamp - t) / (last_timestamp - t[0]) 

    event_frame *= 255
    event_frame = event_frame.astype(dtype=np.uint8)

    stack = np.hstack([frame_bgr, event_frame])
    cv2.imwrite(f"outputs/stack_{idx}.png", stack) 
    
    # cv2.imwrite("frame.png", frame_bgr)

    # input(idx)
    #     #     t, x, y, p = event
    #     # x, y = x.astype(dtype=np.int32), y.astype(dtype=np.int32)

    #     # events = np.hstack([x[..., None], y[..., None], t[..., None], p[..., None]])
    #     # event_labels = segmentation[y, x].astype(dtype=np.uint8)

    #     # write_frame = False
    #     # show_frame = False

    #     # if write_frame or show_frame:
    #     #     ts, xs, ys, ps = event
    #     #     h, w = frame_color.shape[:2]
    #     #     event_bgr = np.zeros((h, w, 3), dtype=np.uint8)
    #     #     for x, y, p in zip(xs, ys, ps):
    #     #         event_bgr[y, x, 0 if p == -1 else 2] = 255
                
    #     #     image_path = image_paths[frame_keys[frame_index]]
    #     #     rgb_image = cv2.imread(image_path)


    # cv2.imshow("Frame", frame)
    # cv2.waitKey(1)

    idx += 1