# Title: Cozmo World Mapping Along a Straight Line
# Authors: Will Churchill, Tao Shi, Sen Zhan
# Date: 5/9/19
# Purpose: For Cozmo Robot to map its world along a straight line, intened for mapping to be used in Monte Carlo Localization

import cozmo
from cozmo.util import degrees, distance_mm, speed_mmps
import os
import sys
import time
import cv2
import copy

import matplotlib.pyplot as plt
from skimage.measure import compare_ssim as ssim

try:
    import numpy as np
except ImportError:
    sys.exit("Cannot import numpy: Do `pip3 install --user numpy` to install")

try:
    from PIL import Image
except ImportError:
    sys.exit("Cannot import from PIL: Do `pip3 install --user Pillow` to install")

# Global list of all pictures taken
pics=[]

# Global list of world map list of pics
picList=[]

# Main Method
# Function: Integrate all methods for fully functioning mapping along a straight line
# Input: Cozmo Robto
# Note: Fully integreated mapping not completed

def main(robot: cozmo.robot.Robot):
    # setting head
   set_head(robot, 0)
   time.sleep(1)
   #make_dir(robot)

   
   for i in range(5):
       # take picture
        robot.camera.image_stream_enabled = True
        robot.add_event_handler(cozmo.world.EvtNewCameraImage, on_new_camera_image)
        time.sleep(0.15)
        robot.camera.image_stream_enabled = False

        #for loop for cozmo to adjust to the light anad redirect itself
        picName = pics[-1]
        while(getLightInCenter(robot,picName)):
            #cozmo take a pic
            robot.camera.image_stream_enabled = True
            robot.add_event_handler(cozmo.world.EvtNewCameraImage, on_new_camera_image)
            time.sleep(0.15)
            robot.camera.image_stream_enabled = False
            picName= pics[-1]
            
        # append to world map list of pics    
        picList.append(picName)
            
        #cozmo move forward
        drive(robot, 70, 3)
        time.sleep(0.15)
 

# Image comparison method
# Function: Compare between two pictures and return the ssim value (1-the same, 0-totally different)
# Note: Not yet integrated in the main method
def compareImages(imageA, imageB):
    image1=cv2.imread(imageA)
    image1=cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    image2=cv2.imread(imageB)
    image2=cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    s = ssim(image1, image2)
    return s
       
# Head Movement Method
# Function: Move Cozmo's head to the desired angle in degrees
# Input: 0 degrees for straight ahead, negative input to lower from 0 and positive input to raise from 0
def set_head(robot: cozmo.robot.Robot, degree):
    robot.set_head_angle(degrees(degree))

# Foward and Backward Movement Method
# Function: Move Cozmo foward or backward for specified time and speed
# Input: speed: negative for backwards, positive for fowards; t: time
def drive(robot: cozmo.robot.Robot, speed, t):
    robot.drive_wheels(speed, speed)
    time.sleep(t)
    
# Picture Taking and Saving Method
# Function: Takes and saves pictures from Cozmo's camera to current directory, also appends new picutures to the global pictures list
# Input: event
def on_new_camera_image(evt, **kwargs):
    pilImage = kwargs['image'].raw_image
    name = "-%d.jpg" % kwargs['image'].image_number
    pics.append(name)
    pilImage.save(name, "JPEG")
    

# Directory Creation Method
# Function: Create directory for saving pictures in (not necessary for functioning of mapping)
# Input: Cozmo Robot
def make_dir(robot: cozmo.robot.Robot):
    if not os.path.exists('pictures'):
        os.makedirs('pictures')
    if not os.path.exists('pictures/'):
        os.makedirs('pictures/')



# Centers Cozmo on Light Source Method
# Function: Rotates Cozmo to horizontally center the light source in the given image
# Input: Cozmo Robot, name of image
def getLightInCenter(robot: cozmo.robot.Robot, imageName):
    image=cv2.imread(imageName
    center=getLightPoint(image, 15)
    width = image.shape[1]
    if center[1]<((width/2)-10):
        robot.turn_in_place(degrees(2))
        return True
    elif center[1]>((width/2)+10):
        robot.turn_in_place(degrees(-2))
        return True
    return False

# Location of Light Source in Image Method
# Function: Determines and returns the location of the light source in the given image
# Input: image: image; circleRadius: approximate size of light source (no need to be specific for our uses)
def getLightPoint(image, circleRadius):
    radius=circleRadius
    orig = image.copy()
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
#use the gaussian blur to make the picture easier to analyze
    gray = cv2.GaussianBlur(gray, (radius, radius), 0)
    (minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(gray)
    image = orig.copy()
    cv2.circle(image, maxLoc, radius, (255, 0, 0), 2)
# display the results of our newly improved method
    return maxLoc



# Cozmo is moved off his charger contacts by default at the start of any program.
# This is because not all motor movement is possible whilst drawing current from
# the charger. In cases where motor movement is not required, such as this example
# we can specify that Cozmo can stay on his charger at the start:
cozmo.robot.Robot.drive_off_charger_on_connect = False

cozmo.run_program(main)