import cv2
import os
import argparse
import sys
import time

from picamera2 import Picamera2
picam2 = Picamera2()

# Define and parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument('--savefolder', help='Folder to save images in',default='Pics')
parser.add_argument('--resolution', help='Desired camera resolution 480p,720p,1080p,4k',default='720p')

args = parser.parse_args()
savefolder = args.savefolder
resolution = args.resolution

frameW = 0
frameH = 0

if (resolution == '480p'):
    frameW = 720
    frameH = 480
elif(resolution == '720p'):
    frameW = 1280
    frameH = 720
elif(resolution == '1080p'):
    frameW = 1920
    frameH = 1080
elif(resolution == '4k'):
    frameW = 3840
    frameH = 2160
else:
    frameW = 1280
    frameH = 720


print(frameW)
print(frameH)

# Create output directory if it doesn't already exist
cwd = os.getcwd()
folderpath = os.path.join(cwd,savefolder)
if not os.path.exists(folderpath):
    os.makedirs(folderpath)

# If images already exist in directory, increment image number so existing images aren't overwritten
# Example: if 'Pics-0.jpg' through 'Pics-10.jpg' already exist, imnum will be incremented to 11
rootname = savefolder
imnum = 1
img_exists = True

while img_exists:
    imname = savefolder + '_' + str(imnum) + '.jpg'
    impath = os.path.join(folderpath, imname)
    if os.path.exists(impath):
        imnum = imnum + 1
    else:
        img_exists = False


picam2.preview_configuration.main.size = (frameW,frameH)
picam2.preview_configuration.main.format = "RGB888"
picam2.preview_configuration.align()
picam2.configure("preview")
picam2.start()


while True:
    frame= picam2.capture_array()
    cv2.imshow("Camera", frame)
    key = cv2.waitKey(1)
    if key==ord('q'):
        break
    elif key == ord('p'):
        #Take a picture!
        filename = savefolder + '_' + str(imnum) + '.jpg'
        savepath = os.path.join(folderpath, filename)
        cv2.imwrite(savepath, frame)
        print('Picture taken and saved as %s' % filename)
        imnum = imnum + 1
cv2.destroyAllWindows()