import argparse
import pathlib
import zipfile

import cv2
import numpy as np

LABEL_VALUES = {
    100: 'RED',
    200: 'ORANGE',
    300: 'CYAN',
    400: 'BLUE',
    # ----------------
    65535: 'NO_LABEL',  # 2**16-1
}

COLORS = {
    'CYAN':     (255, 191, 0),
    'BLUE':     (255, 128, 0),
    'ORANGE':   (0, 128, 255),
    'RED':      (0, 0, 255),
    # ----------------
    'NO_LABEL': (127,127,127),
}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='DVS Zip-CSV reader example')
    parser.add_argument("--input", type=pathlib.Path, required=True, help="Path to input zip file")
    args = parser.parse_args()
    
    if not args.input.is_file():
        raise RuntimeError("input invalid")

    archive = zipfile.ZipFile(args.input, 'r')
    eventDataFiles = sorted([fpath for fpath in archive.namelist() if fpath.endswith(".csv")])
    
    for eventCsv in eventDataFiles:
        print("Processing", eventCsv)
    
        polarity_frame = np.zeros((640,768,3), dtype=np.uint8)
        label_frame = np.zeros((640,768,3), dtype=np.uint8)
        with archive.open(eventCsv, 'r') as fd:
            for line in fd:
                line = line.decode('utf-8').strip() 

                # x y t[ms] polarity label
                x, y, t, p, l = line.split(' ')
                x, y, t, p, l = int(x), int(y), float(t), int(p), int(l)
                
                if p == -1:
                    polarity_frame[y, x] = (0, 0, 255)
                elif p == 0:
                    polarity_frame[y, x] = (255, 0, 0)
                elif p == 1:
                    polarity_frame[y, x] = (0, 255, 0)
                else:
                    raise ValueError("unkown polarity")
                    
                label_frame[y, x] = COLORS[LABEL_VALUES[l]]
                    
            cv2.imshow("polarity", polarity_frame)
            cv2.imshow("label", label_frame)
            cv2.waitKey(1)