#!/usr/bin/env python3

import pyrealsense2 as rs
import numpy as np
import cv2
import requests
import os
import json
from tempfile import TemporaryFile
from datetime import datetime


class DepthCapture:
    def __init__(self, fps=30, window_name="RealSense"):
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.pipeline_wrapper = rs.pipeline_wrapper(self.pipeline)
        self.pipeline_profile = self.config.resolve(self.pipeline_wrapper)
        self.device = self.pipeline_profile.get_device()
        self.queue = rs.frame_queue(50, keep_frames=True)
        self.depth_scale = self.device.first_depth_sensor().get_depth_scale()

        self.config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, fps)
        self.config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, fps)
        
        self.window_name = window_name
        self.flash = False
        self.lock = False

    def touchcb(self, event, *_):
        if event == cv2.EVENT_LBUTTONDOWN and self.lock == False:
            #Set flag so main loop can handle stuff
            self.flash = True
    
    def run(self):
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setWindowProperty(self.window_name, cv2.WND_PROP_FULLSCREEN,cv2.WINDOW_FULLSCREEN)
        cv2.setMouseCallback(self.window_name, self.touchcb)
        
        self.pipeline.start(self.config,self.queue)
        sess = requests.Session()
        try:
            while True:
                frames = self.queue.wait_for_frame().as_frameset()
                color_frame = frames.get_color_frame()
                color_image = np.asanyarray(color_frame.get_data())
                if self.flash:
                    self.lock = True
                    self.flash = False
                    cv2.imshow(self.window_name, np.full_like(color_image, 255))
                    cv2.waitKey(1)
                    depth_frame = frames.get_depth_frame()
                    depth_image = np.asanyarray(depth_frame.get_data())
                    intrinsics = depth_frame.profile.as_video_stream_profile().get_intrinsics()
                    with TemporaryFile() as tmpfile:
                        np.savez(tmpfile, depth=depth_image, color=color_image)
                        tmpfile.seek(0)
                        sess.post('UPLOAD URL REMOVED',
                            headers={'Authorization':f'Bearer {os.environ["RETE_SECRET"]}'},
                            data={'data':json.dumps({"depth_scale":self.depth_scale,
                                                     "dimensions":(intrinsics.height, intrinsics.width),
                                                     "focal_lengths":(intrinsics.fx, intrinsics.fy),
                                                     "principal_point":(intrinsics.ppx, intrinsics.ppy)})},
                            files={'file0':(f'{datetime.now()}.npz', tmpfile)})
                    cv2.waitKey(50)
                    self.lock = False
                else:
                    cv2.imshow(self.window_name, color_image)
                    cv2.waitKey(1)
        finally:
            self.pipeline.stop()
            sess.close()
    

cap = DepthCapture()
cap.run()