import numpy as np
import cv2

# Identify pixels above the threshold
# Threshold of RGB > 160 does a nice job of identifying ground pixels only
def color_thresh(img, rgb_thresh=(160, 160, 160)):
    # Create an array of zeros same xy size as img, but single channel
    color_select = np.zeros_like(img[:,:,0])
    # Require that each pixel be above all three threshold values in RGB
    # above_thresh will now contain a boolean array with "True"
    # where threshold was met
    above_thresh = (img[:,:,0] > rgb_thresh[0]) \
                 & (img[:,:,1] > rgb_thresh[1]) \
                 & (img[:,:,2] > rgb_thresh[2])
    # Index the array of zeros with the boolean array and set to 1
    color_select[above_thresh] = 1
    # Return the binary image
    return color_select


def obstacle_threshold(img, rgb_threshold=(160,160,160)):
    """
        color filtering for non-navigable terrain.
    """
    # create an array of zeros the same xy size as the img and single channel
    obstacle_selection = np.zeros_like(img[:,:,0])
    # Require that each pixel is below the three threshold values in RGB
    # below_threshold will now contain all beneath threshold values
    below_threshold = (img[:,:,0] < rgb_threshold[0])&\
                      (img[:,:,0] < rgb_threshold[1])&\
                      (img[:,:,0] < rgb_threshold[2])
    obstacle_selection[below_threshold] = 1
    return obstacle_selection

def rock_threshold(img, rgb_threshold=None):
    """
        Color filtering for mapping out rock samples, note:
        rgb_thresholds=None but in the function two thresholds are set.
    """
    rgb_threshold_upper=(255, 200, 50)
    rgb_threshold_lower=(150, 100, 0)
    # Place the code to color select the rock samples the rocks are yellow
    rock_selection = np.zeros_like(img[:,:,0])
    # this is a specific color-range so it needs two values upper and lower limits
    rock_threshold = (img[:,:,0] < rgb_threshold_upper[0])&\
                     (img[:,:,0] < rgb_threshold_upper[1])&\
                     (img[:,:,0] < rgb_threshold_upper[2])&\
                     (img[:,:,0] > rgb_threshold_lower[0])&\
                     (img[:,:,0] > rgb_threshold_lower[1])&\
                     (img[:,:,0] > rgb_threshold_lower[2])


    rock_selection[rock_threshold] = 1
    return rock_selection

# from the course
def find_rocks(img, levels=(110, 110, 50)):
    rockpix = ((img[:,:,0] > levels[0]) &\
               (img[:,:,1] > levels[1]) &\
               (img[:,:,2] < levels[2]))
    color_select = np.zeros_like(img[:,:,0])
    color_select[rockpix] = 1
    return color_select

# Define a function to convert from image coords to rover coords
def rover_coords(binary_img):
    # Identify nonzero pixels
    ypos, xpos = binary_img.nonzero()
    # Calculate pixel positions with reference to the rover position being at the
    # center bottom of the image.
    x_pixel = -(ypos - binary_img.shape[0]).astype(np.float)
    y_pixel = -(xpos - binary_img.shape[1]/2 ).astype(np.float)
    return x_pixel, y_pixel


# Define a function to convert to radial coords in rover space
def to_polar_coords(x_pixel, y_pixel):
    # Convert (x_pixel, y_pixel) to (distance, angle)
    # in polar coordinates in rover space
    # Calculate distance to each pixel
    dist = np.sqrt(x_pixel**2 + y_pixel**2)
    # Calculate angle away from vertical for each pixel
    angles = np.arctan2(y_pixel, x_pixel)
    return dist, angles

# Define a function to map rover space pixels to world space
def rotate_pix(xpix, ypix, yaw):
    # Convert yaw to radians
    yaw_rad = yaw * np.pi / 180
    xpix_rotated = (xpix * np.cos(yaw_rad)) - (ypix * np.sin(yaw_rad))

    ypix_rotated = (xpix * np.sin(yaw_rad)) + (ypix * np.cos(yaw_rad))
    # Return the result
    return xpix_rotated, ypix_rotated

def translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale):
    # Apply a scaling and a translation
    xpix_translated = (xpix_rot / scale) + xpos
    ypix_translated = (ypix_rot / scale) + ypos
    # Return the result
    return xpix_translated, ypix_translated


# Define a function to apply rotation and translation (and clipping)
# Once you define the two functions above this function should work
def pix_to_world(xpix, ypix, xpos, ypos, yaw, world_size, scale):
    # Apply rotation
    xpix_rot, ypix_rot = rotate_pix(xpix, ypix, yaw)
    # Apply translation
    xpix_tran, ypix_tran = translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale)
    # Perform rotation, translation and clipping all at once
    x_pix_world = np.clip(np.int_(xpix_tran), 0, world_size - 1)
    y_pix_world = np.clip(np.int_(ypix_tran), 0, world_size - 1)
    # Return the result
    return x_pix_world, y_pix_world

# Define a function to perform a perspective transform
def perspect_transform(img, src, dst):

    M = cv2.getPerspectiveTransform(src, dst)
    warped = cv2.warpPerspective(img, M, (img.shape[1], img.shape[0]))
    # from the walkthrough video (because color thresholding the obstacles was not working out well)
    mask = cv2.warpPerspective(np.ones_like(img[:,:,0]), M, (img.shape[1], img.shape[0]))
    # keep same size as input image
    return warped, mask


# Apply the above functions in succession and update
# the Rover state accordingly
def perception_step(Rover):
    # Perform perception steps to update Rover()
    # TODO:
    # NOTE: camera image is coming to you in Rover.img
    # 1) Define source and destination points for perspective transform
    # Use Rover.img as image
    bottom_offset = 6
    dst_size = 5
    source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
    destination = np.float32([
                  [Rover.img.shape[1]/2 - dst_size, Rover.img.shape[0] - bottom_offset],
                  [Rover.img.shape[1]/2 + dst_size, Rover.img.shape[0] - bottom_offset],
                  [Rover.img.shape[1]/2 + dst_size, Rover.img.shape[0] - 2*dst_size - bottom_offset],
                  [Rover.img.shape[1]/2 - dst_size, Rover.img.shape[0] - 2*dst_size - bottom_offset],
                  ])
    # 2) Apply perspective transform
    warped, mask = perspect_transform(Rover.img, source, destination)
    # 3) Apply color threshold to identify navigable
    #    terrain/obstacles/rock samples
    navigable_terrain = color_thresh(warped)
    obstacles = np.absolute(np.float32(navigable_terrain) - 1) * mask
    # obstacles = obstacle_threshold(warped)
    rock_map = find_rocks(warped, levels=(110, 110, 50))
    scale = 2*dst_size
    world_size = Rover.worldmap.shape[0]
    if rock_map.any():
        rock_x, rock_y = rover_coords(rock_map)
        rock_x_world, rock_y_world = \
        pix_to_world(rock_x, rock_y, Rover.pos[0],\
                     Rover.pos[1], Rover.yaw, world_size, scale)
        rock_dist, rock_ang = to_polar_coords(rock_x, rock_y)
        rock_idx = np.argmin(rock_dist)
        rock_xcen = rock_x_world[rock_idx]
        rock_ycen = rock_y_world[rock_idx]
        Rover.worldmap[rock_ycen, rock_xcen, 1] = 255
        Rover.vision_image[:,:,1] = rock_map * 255
    else:
        Rover.vision_image[:,:,1] = 0
    # 4) Update Rover.vision_image
    # (this will be displayed on left side of screen)
    # obstacle color-thresholded binary image
    # Rover.vision_image[:,:,1] = rocks*255 # white
    # rock_sample color-thresholded binary image
    Rover.vision_image[:,:,0] = obstacles*255 # Red
    Rover.vision_image[:,:,2] = navigable_terrain*255 # Blue
    # navigable terrain color-thresholded binary image
    # 5) Convert map image pixel values to rover-centric coords
    rover_x, rover_y = rover_coords(navigable_terrain)
    rover_x_obstacle, rover_y_obstacle = rover_coords(obstacles)
    # rover_x_rock, rover_y_rock = rover_coords(rocks)

    # 6) Convert rover-centric pixel values to world coordinates
    obstacle_x_world, obstacle_y_world = pix_to_world(rover_x_obstacle,
                                                      rover_y_obstacle,
                                                      Rover.pos[0],
                                                      Rover.pos[1],
                                                      Rover.yaw,
                                                      world_size,
                                                      scale)

    # rock_x_world, rock_y_world = pix_to_world(rover_x_rock,
    #                                           rover_y_rock,
    #                                           Rover.pos[0],
    #                                           Rover.pos[1],
    #                                           Rover.yaw,
    #                                           world_size,
    #                                           scale)

    navigable_x_world, navigable_y_world = pix_to_world(rover_x,
                                                        rover_y,
                                                        Rover.pos[0],
                                                        Rover.pos[1],
                                                        Rover.yaw,
                                                        world_size,
                                                        scale)

    # 7) Update Rover worldmap (to be displayed on right side of screen)
    Rover.worldmap[obstacle_y_world, obstacle_x_world, 0] += 1 # Red
    # Rover.worldmap[rock_y_world, rock_x_world, 1] = 255
    Rover.worldmap[navigable_y_world, navigable_x_world, 2] += 10 # Blue

    # 8) Convert rover-centric pixel positions to polar coordinates
    # these are the angles and distances from navigable pixels
    distances, angles = to_polar_coords(rover_x, rover_y)

    # Update Rover pixel distances and angles
    Rover.nav_dists = distances
    Rover.nav_angles = angles
    # Do we want distances and angles for obtacles and rocks?
    obstacle_distances, obstacle_angles = to_polar_coords(rover_x_obstacle, rover_y_obstacle)
    # rock_distances, rock_angles = to_polar_coords(rover_x_rock, rover_y_rock)
    Rover.obstacle_dists = obstacle_distances
    Rover.obstacle_angles = obstacle_angles
    # Rover.rock_dists = rock_distances
    # Rover.rock_angles = rock_angles
    return Rover
