import serial
import time
import struct

# Small example showing how to get an image from the robot, convert it from RGB565 to RGB888 and finally save to bmp.
# Beware that the total image size must not exceed 4056 bytes.
# The default image settings are: rgb565, 40x40, 4x subsampling, at the center (160,80)

IMAGE_WIDTH = 40
IMAGE_HEIGHT = 40

brg888 = bytearray([0] * (IMAGE_WIDTH*IMAGE_HEIGHT*3)) # After conversion the image has 3 bytes per pixel

def save_bmp_image(filename, image, width, height):
	filesize = 54 + 3 * width * height
	print("filesize = " + str(filesize))
	bmpfileheader = bytearray(14)
	bmpfileheader[0] = ord('B')
	bmpfileheader[1] = ord('M')
	bmpfileheader[2] = filesize&0xFF
	bmpfileheader[3] = (filesize >> 8)&0xFF
	bmpfileheader[4] = (filesize >> 16)&0xFF
	bmpfileheader[5] = (filesize >> 24)&0xFF
	bmpfileheader[6] = 0
	bmpfileheader[7] = 0
	bmpfileheader[8] = 0
	bmpfileheader[9] = 0
	bmpfileheader[10] = 54
	bmpfileheader[11] = 0
	bmpfileheader[12] = 0
	bmpfileheader[13] = 0
	
	bmpinfoheader = bytearray(40)
	bmpinfoheader[0] = 40
	bmpinfoheader[1] = 0
	bmpinfoheader[2] = 0
	bmpinfoheader[3] = 0
	bmpinfoheader[4] = width&0xFF
	bmpinfoheader[5] = (width >> 8)&0xFF
	bmpinfoheader[6] = (width >> 16)&0xFF
	bmpinfoheader[7] = (width >> 24)&0xFF
	bmpinfoheader[8] = (height)&0xFF
	bmpinfoheader[9] = (height >> 8)&0xFF
	bmpinfoheader[10] = (height >> 16)&0xFF
	bmpinfoheader[11] = (height >> 24)&0xFF
	bmpinfoheader[12] = 1
	bmpinfoheader[13] = 0
	bmpinfoheader[14] = 24
	bmpinfoheader[15] = 0

	bmppad = bytearray(3)

	with open(filename,'wb') as file:
		file.write(bmpfileheader)
		file.write(bmpinfoheader)
		for i in range(height):
			file.write(image[(width * (height - i - 1) * 3):(width * (height - i - 1) * 3)+(3 * width)])
			file.write(bmppad[0:((4 - (width * 3) % 4) % 4)])			

def grey_to_brg888(grey, brg888, width, height):
	counter = 0
	for j in range(height):
		for i in range(width):
			index = 3 * (i + j * width)
			brg888[index] = grey[counter]
			brg888[index + 1] = grey[counter]
			brg888[index + 2] = grey[counter]
			counter += 1
			
def rgb565_to_brg888(rgb565, brg888, width, height):
	counter = 0
	for j in range(height):
		for i in range(width):
			index = 3 * (i + j * width)
			#print("rgb565: " + str(rgb565[counter]) + ", " + str(rgb565[counter+1]))
			red = rgb565[counter]&0xF8
			green = ((rgb565[counter]&0x07) << 5) & 0xFF
			counter += 1
			green += ((rgb565[counter]&0xE0) >> 3)
			blue = ((rgb565[counter]&0x1F) << 3) & 0xFF
			counter += 1
			#print("r=" + str(red) + ", g=" + str(green) + ", b=" + str(blue))
			brg888[index] = blue
			brg888[index + 1] = green
			brg888[index + 2] = red

ser = serial.Serial('COM59', 115200, timeout=0)

message = bytearray("J,1," + str(IMAGE_WIDTH) + "," + str(IMAGE_HEIGHT) + ",4,160,80\n", 'utf-8')
ser.write(message)
reply = ser.read()
while len(reply) < 3:
	reply += ser.read()
print("Camera setting ok: " + str(reply))

message = struct.pack(">bb", - ord('I'), 0)
ser.write(message)
reply = ser.read()
while len(reply) < (IMAGE_WIDTH*IMAGE_HEIGHT*2+3): # width*height*bytes_per_pixel + header
	reply += ser.read()

ser.close()
rgb565_to_brg888(reply[3:], brg888, IMAGE_WIDTH, IMAGE_HEIGHT)
save_bmp_image("image.bmp", brg888, IMAGE_WIDTH, IMAGE_HEIGHT)