#! /usr/bin/python3

import rospy, sys, math, cv2, numpy as np, std_msgs.msg
from sensor_msgs.msg import Image
from darknet_ros_msgs.msg import BoundingBox
from std_msgs.msg import Header
import os.path, getpass
from cv_bridge import CvBridge, CvBridgeError

# gets processed images from darknet image topic and doesnt save them
class ImageProcessor():
    def __init__(self):
        self.bridge = CvBridge()
        self.imageSubscriber = rospy.Subscriber("webcam/image_raw", Image, self.imageCallback)
        self.boxSubscriber = rospy.Subscriber("darknet_ros/bounding_boxes", BoundingBox, self.boxCallback)
        self.params = None
        self.image = None
        self.seq = None

        # stuff that needs to be parameters
        self.objectsOfInterest = ["person", "backpack", "car"] # things we report
        self.reportConfidence = 0.80 # confidence required for report
        self.reportFrames = 10 # number of frames at high confidence before reporting
        self.saveDirectory = "/tmp/darknetMemories/"

    def boxCallback(self, box):
        # check confidence levels and save image when sure of object
        saveFlag = True
        if saveFlag:
            if os.path.isdir(self.saveDirectory):
                cv2.imwrite(self.saveDirectory + "camera_image_{seq}.jpeg".format(seq = self.seq), self.image)
                rospy.loginfo("saving image to folder: " + self.saveDirectory + "-> camera_image_{seq}.jpeg".format(seq = self.seq))
                # cv2.waitKey(10)
            else:
                rospy.loginfo("image save process failed!")

    def imageCallback(self, frame):
        # load image from topic for later use
        try:
            self.image = self.bridge.imgmsg_to_cv2(frame, "bgr8")
            self.seq = frame.header.seq
        except CvBridgeError as err:
            rospy.loginfo("there was error!")
            rospy.loginfo(err)

def main():
    rospy.init_node("darknetImageProcessor")
    iv = ImageProcessor()

    rospy.spin()
    # while not rospy.is_shutdown():
        # if iv.image is not None:
        #     if np.mod(iv.seq, 32) == 0:
        #         directory = "/home/" + getpass.getuser() + "/" + folderName
        #         if os.path.isdir(directory):
        #             if windowFlag:
        #                 cv2.imshow("image window", iv.image)
        #                 rospy.loginfo("publishing image to window seq: " + iv.seq)
        #             if saveFlag:
        #                 cv2.imwrite(directory + "/camera_image_{seq}.jpeg".format(seq = iv.seq), iv.image)
        #                 rospy.loginfo("saving image to folder: " + directory + "-> camera_image_{seq}.jpeg".format(seq = iv.seq) + " seq: " + str(iv.seq))
        #             cv2.waitKey(10)
        #         else:
        #             rospy.loginfo("camera dir missing!")

    # cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
