import cv2
import json
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from metavision_core.event_io import EventsIterator


def visualization_event_streams(p_list, t_list, x_list, y_list, save_path=None):
    # 事件流的3D表示
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.xaxis.set_pane_color((1.0, 1.0, 1.0))  # X轴背景颜色为透明
    ax.yaxis.set_pane_color((1.0, 1.0, 1.0))  # Y轴背景颜色为透明
    ax.zaxis.set_pane_color((1.0, 1.0, 1.0))  # Z轴背景颜色为透明
    ax.grid(False)


    new_t1 = [(t-17061000)/100000 for t in t_list]                  # 调整时间戳

    p1 = np.array(p_list)
    x1 = np.array(x_list)
    y1 = np.array(y_list)
    t1 = np.array(new_t1)
    colors = np.where(p1 == 1, 'red', 'blue')
    ax.scatter(x1, y1, t1, c=colors, marker='o',s=0.5)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('t(s)')
    ax.set_xlim(0, 1280)
    ax.set_ylim(0, 720)
    # print(t1)

    if save_path is not None:
        plt.savefig(str(save_path), bbox_inches='tight')
    plt.close()



# 将点生成2Dmap, TARGET_SHAPE = (346, 260)，(720, 1280, 3)
def generate_2D_map(height, width, dict_file, save_map_path, save_stream_path):
    map = np.ones((height, width, 3))
    map = map * 255

    p_list = dict_file['p']
    t_list = dict_file['t']
    x_list = dict_file['x']
    y_list = dict_file['y']

    # 生成斜45度事件流图
    visualization_event_streams(p_list, t_list, x_list, y_list, save_stream_path)

    # 生成2dmap
    for polarity, xx, yy in zip(p_list, x_list, y_list):
        yy, xx = int(yy), int(xx)

        # 正事件为红色，负事件为蓝色，numpy：BGR
        if polarity == 1:
            map[yy][xx][0] = 0
            map[yy][xx][1] = 0
            map[yy][xx][2] = 255
        elif polarity == 0:
            map[yy][xx][0] = 255
            map[yy][xx][1] = 0
            map[yy][xx][2] = 0
        else:
            raise BaseException(f"极性错误！({xx},{yy}) {polarity} {save_map_path}")

    cv2.imwrite(str(save_map_path), map)






if __name__ == "__main__":
    """ Main """

    event_file_path = 'UAVforPPT.raw'  # 输入的 .raw 文件路径

    save_json_dir = Path("UAVforPPT/json")
    save_json_dir.mkdir(exist_ok=True, parents=True)
    save_event_frame = Path("UAVforPPT/event_frame")
    save_event_frame.mkdir(exist_ok=True, parents=True)
    save_event_stream = Path("UAVforPPT/event_stream")
    save_event_stream.mkdir(exist_ok=True, parents=True)

    # Events iterator on Camera or event file
    duration_time_us = 3000    # us
    mv_iterator = EventsIterator(input_path=event_file_path, delta_t=duration_time_us)
    height, width = mv_iterator.get_size()  # Camera Geometry

    idx = 0
    global_counter = 0  # This will track how many events we processed
    global_max_t = 0  # This will track the highest timestamp we processed
    dict_file = dict(p=[], t=[], x=[], y=[])

    # Process events
    for evs in mv_iterator:
        print("----- New event buffer! -----")
        if evs.size == 0:
            print("The current event buffer is empty.\n")
        else:
            min_t = evs['t'][0]   # Get the timestamp of the first event of this callback
            max_t = evs['t'][-1]  # Get the timestamp of the last event of this callback
            global_max_t = max_t  # Events are ordered by timestamp, so the current last event has the highest timestamp

            counter = evs.size  # Local counter
            global_counter += counter  # Increase global counter

            print(f"There were {counter} events in this event buffer.")
            print(f"There were {global_counter} total events up to now.")
            print(f"The current event buffer included events from {min_t} to {max_t} microseconds.")
            print("----- End of the event buffer! -----\n")


            if min_t < 17061000:
                continue
            if max_t > 17363997:
                break

            dict_file['t'] = [int(i) for i in evs['t']]
            dict_file['p'] = [int(i) for i in evs['p']]
            dict_file['x'] = [int(i) for i in evs['x']]
            dict_file['y'] = [int(i) for i in evs['y']]

            # 保存json文件
            file_name = f"{idx}_{min_t}-{max_t}_{duration_time_us}us"
            save_json = save_json_dir / f'{file_name}.json'
            # if save_json.exists():  continue
            with open(save_json, 'w') as handle:
                json.dump(dict_file, handle)

            # 生成可视化图
            generate_2D_map(height, width, dict_file, save_event_frame / f'{file_name}.png', save_event_stream / f'{file_name}.png')

            dict_file['p'].clear()
            dict_file['t'].clear()
            dict_file['x'].clear()
            dict_file['y'].clear()
            idx += 1


            # break


    # Print the global statistics
    duration_seconds = global_max_t / 1.0e6
    print(f"There were {global_counter} events in total.")
    print(f"The total duration was {duration_seconds:.2f} seconds.")
    if duration_seconds >= 1:  # No need to print this statistics if the total duration was too short
        print(f"There were {global_counter / duration_seconds :.2f} events per second on average.")