#!/usr/bin/env python
import glob
import time
import rospy
from std_msgs.msg import String
import numpy as np
import cv2 as cv

from include.game_lib import Player
import include.calibration_lib as calib
import include.yaml_data_handling as yaml_handle

# Debug mode, when True, show a lot of useful messages
DEBUG = rospy.get_param("robocheckers/debug")

# Pieces publisher node initialization
rospy.init_node('game_detection')
pub = rospy.Publisher('/robocheckers/pieces', String, queue_size=10)

# Decide whether to use live image or stored images
USE_LIVE_IMAGES = True

# Offset in pixels to extend chessboard cropped image (default=16)
CROP_OFFSET = 16

## Defining various filters for further use
# Green filter
GREEN = ((30, 30, 10), (100, 255, 200))
# Orange filter
ORANGE = ((0, 100, 5), (25, 255, 255))
# Blue filter 
BLUE = ((100, 100, 5), (135, 255, 255))


def mask2Binary(img, thresholds, minColor):
    # Convert to HSV colorspace
    blur = cv.GaussianBlur(img, (3, 3), 1.2)
    hsv = cv.cvtColor(blur, cv.COLOR_BGR2HSV)

    # Extract upper and lower values from thresholds
    if len(thresholds) == 2:
        lower, upper = thresholds
        mask = cv.inRange(hsv, lower, upper)
    # Other
    else:
        print("Game Detection ERROR: Invalid format!")
        return -1

    # Apply color mask and convert to binary
    color = cv.bitwise_and(hsv, hsv, mask=mask)
    value = color[:, :, 2]

    thresh = cv.threshold(value, minColor, 255, cv.THRESH_BINARY)[1]

    # Apply morphology open to remove noise and fill holes in marks
    kernel = np.ones((7, 7), np.uint8)
    thresh = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel)
    kernel = np.ones((3, 3), np.uint8)
    thresh = cv.morphologyEx(thresh, cv.MORPH_CLOSE, kernel)

    return thresh


def merge_images(img1, img2):
    """
    Merges 2 binary images into one and covers the holes
    :param img1: First binary image
    :param img2: Second binary image
    :return: Merged and fixed image
    """
    # Merge images
    img = img1 | img2

    # Apply morphology close fill holes in blobs
    kernel = np.ones((3, 3), np.uint8)
    img = cv.morphologyEx(img, cv.MORPH_CLOSE, kernel)

    return img


def crop_by_offset(img, offset, center=None):
    """
    Crops the image by give offset from all four sides
    :param img: Image to be cropped
    :param offset: Offset to crop by
    :return: Cropped image
    """
    if offset:
        img = img[offset:-offset, offset:-offset]

    if center:
        tmp = list(center)
        tmp[0], tmp[1] = center[0]-offset, center[1]-offset
        center = tuple(tmp)
        return img, center
    return img


def resize_chessboard(frame, pieces_a, pieces_b, camera_center=None):
    """
    Resize the game chessboard and normalize all pieces positions by new size
    :param frame: Image to be resized
    :param shape: Shape to which the chessboard will be resized to
    :param pieces_a: Pieces of the first player
    :param pieces_b: Pieces of the second player
    :param camera_center: Centerpoint of view of camera
    :return result: Resized image
    :return pieces_a: Normalized pieces from player one
    :return pieces_v: Normalized pieces from player two
    :return camera_center: Normalized camera center
    """
    # Saving the old shape values into variable
    old_shape = frame.shape[:2]

    # Getting new shape from dimension_data.yml and creating a tuple from it
    chessboard = yaml_handle.load_dimensions_data(yaml_handle.RESIZE_TO_PX)
    chessboard = chessboard.get('chessboard')
    shape = (chessboard, chessboard)

    # Reshape the chessboard
    result = cv.resize(frame, shape, cv.INTER_LANCZOS4)

    for i, piece in enumerate(pieces_a):
        pieces_a[i].pos = list(pieces_a[i].pos)
        pieces_a[i].pos[0] = piece.pos[0]/old_shape[0]*shape[0]
        pieces_a[i].pos[1] = piece.pos[1]/old_shape[1]*shape[1]
        pieces_a[i].pos = tuple(pieces_a[i].pos)

    for i, piece in enumerate(pieces_b):
        pieces_b[i].pos = list(pieces_b[i].pos)
        pieces_b[i].pos[0] = piece.pos[0]/old_shape[0]*shape[0]
        pieces_b[i].pos[1] = piece.pos[1]/old_shape[1]*shape[1]
        pieces_b[i].pos = tuple(pieces_b[i].pos)

    # If not counting with camera center, return
    if camera_center is None:
        return result, pieces_a, pieces_b

    # Preparing for reshaping the camera_center
    width, height = frame.shape[:2]
    camera_center = list(camera_center)
    camera_center[0] = int(round(float(camera_center[0])/width*shape[0]))
    camera_center[1] = int(round(float(camera_center[1])/height*shape[1]))
    camera_center = tuple(camera_center)

    return result, pieces_a, pieces_b, camera_center


def blob_detect(img, minArea):
    # Setup SimpleBlobDetector parameters
    params = cv.SimpleBlobDetector_Params()

    # Change thresholds
    params.minThreshold = 239
    params.thresholdStep = 8
    params.maxThreshold = 255

    # Filter by Area
    params.filterByArea = True
    params.minArea = minArea

    # Filter by Color
    params.filterByColor = True
    params.blobColor = 255

    # Filter by Convexity
    params.filterByConvexity = True
    params.minConvexity = 0.1

    # Filter by Inertia
    params.filterByInertia = False

    detector = cv.SimpleBlobDetector_create(params)

    # Detect blobs
    keypoints = detector.detect(img)

    return keypoints


def blob_draw(img, keypoints, color_select):
    # Color selection
    if color_select == "green":
        color = (0, 255, 0)
    elif color_select == "orange":
        color = (0, 127, 255)
    elif color_select == "blue":
        color = (255, 127, 0)
    else:
        return None

    return_image = cv.drawKeypoints(img, keypoints, np.array([]), color,
                                    cv.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

    return return_image


def piece_draw(img, pieces, color_select):
    # Color selection
    if color_select == "orange":
        color = (0, 127, 255)
    elif color_select == "blue":
        color = (255, 127, 0)
    else:
        return False

    # Set green color for mark visualation
    mark_color = (0, 255, 0)
    
    # Getting new shape from dimension_data.yml and creating a tuple from it
    real_boardsize = yaml_handle.load_dimensions_data(yaml_handle.REAL_MM)
    real_boardsize = real_boardsize.get('chessboard')

    # Calculating radii for pieces and marks
    avg_len = sum(img.shape[:2])/2
    radius = int(round(((16.0/real_boardsize)*avg_len), 0))
    mark_radius = int(round(((7.5/real_boardsize)*avg_len), 0))

    # Draw detected blobs as pieces
    for piece in pieces:
        cv.circle(img, (int(piece.pos[0]), int(piece.pos[1])), radius, color, 1)
        if piece.rank == "king":
            cv.circle(img, (int(piece.pos[0]), int(piece.pos[1])), mark_radius, mark_color, 1)

    return True


def main():
    if USE_LIVE_IMAGES:
        # Camera selection
        cap = cv.VideoCapture(rospy.get_param("robocheckers/camera_id"))
    else:
        # Load images
        images = sorted(glob.glob(yaml_handle.CURRENT_DIR + '/test/images/*.jpg'))
        frame_index = 0

    # Load calibration data
    mtx, dist = yaml_handle.load_camera_params()
    chessboard_center, angle, apexes = yaml_handle.load_chessboard_params()

    # Player class creation
    Blue = Player()
    Orange = Player()

    while True:
        if USE_LIVE_IMAGES:
            # Capture frame-by-frame
            ret, frame = cap.read()

            # If no image is capped, return
            if not ret:
                print "Game Detection ERROR: Camera not found!"
                rospy.sleep(5)
                return
        else:
            # Cycle through the globbed images
            if frame_index >= len(images):
                frame_index = 0
            if frame_index < 0:
                frame_index = len(images)-1
            frame = cv.imread(images[frame_index])

        #y, x, _ = frame.shape
        #cv.circle(frame, (x/2, y/2), radius=6, color=(0, 255, 0), thickness=-1)

        # Calibrate the frame
        dst, camera_center = calib.calibrate_frame(frame, mtx, dist)
        # Rotate by angle plus 180 degree
        dst, camera_center = calib.rotate_frame(dst, chessboard_center, angle+np.pi,
                                                       camera_center)
        # Crop the chessboard
        chessboard, camera_center = calib.adjust_frame(dst.copy(), apexes, camera_center,
                                                              CROP_OFFSET)

        # Filtering colors
        green_frame = mask2Binary(chessboard, GREEN, 10)
        orange_frame = mask2Binary(chessboard, ORANGE, 30)
        blue_frame = mask2Binary(chessboard, BLUE, 10)

        # Merge the pieces images with marks images
        orange_blob_frame = merge_images(orange_frame, green_frame)
        blue_blob_frame = merge_images(blue_frame, green_frame)

        # Find blobs
        green_blobs = blob_detect(green_frame, 200)
        orange_blobs = blob_detect(orange_blob_frame, 1000)
        blue_blobs = blob_detect(blue_blob_frame, 1000)

        green_blob_frame = blob_draw(green_frame, green_blobs, "green")
        orange_blob_frame = blob_draw(orange_blob_frame, orange_blobs, "orange")
        blue_blob_frame = blob_draw(blue_blob_frame, blue_blobs, "blue")

        # Orange and Blue pieces information storing into instances of player classes
        Orange.get_piece_data(orange_blobs, green_blobs, CROP_OFFSET)
        Blue.get_piece_data(blue_blobs, green_blobs, CROP_OFFSET)

        # Chessboard visualization
        chessboard = cv.cvtColor(chessboard, cv.COLOR_BGR2GRAY)
        #chessboard = cv.threshold(chessboard, 100, 255, cv.THRESH_BINARY)[1]
        chessboard = cv.cvtColor(chessboard, cv.COLOR_GRAY2BGR)

        # Crop image by offset and resize it to 512x512 px
        chessboard, camera_center = crop_by_offset(chessboard, CROP_OFFSET, camera_center)
        chessboard, Orange.pieces, Blue.pieces, camera_center = resize_chessboard(chessboard,
                                                                                  Orange.pieces,
                                                                                  Blue.pieces,
                                                                                  camera_center)
        cv.circle(chessboard, camera_center, radius=3, color=(0, 0, 255), thickness=-1)

        # Calculate the distances between pieces and camera frame center
        Orange.calc_dists(camera_center)
        Blue.calc_dists(camera_center)

        Orange.normalize_pieces()
        Blue.normalize_pieces()

        if DEBUG:
            print "===== Orange Player ====="
            print Orange

            print "===== Blue Player ====="
            print Blue

        Orange.publish(pub, "orange")
        Blue.publish(pub, "blue")
        if DEBUG: print "Camera center:", camera_center

        # Draw all pieces
        piece_draw(chessboard, Orange.pieces, "orange")
        piece_draw(chessboard, Blue.pieces, "blue")

        cv.imshow("cropped", chessboard)
        #all_frame = np.hstack(((np.hstack((green_frame, orange_frame))), blue_frame))
        if DEBUG:
            all_frame = np.hstack(((np.hstack((green_blob_frame, orange_blob_frame))), blue_blob_frame))
            cv.imshow('Frames: Green, Orange, Blue', all_frame)

            cv.imshow('Frame', dst)

        if USE_LIVE_IMAGES:
            if cv.waitKey(1) & 0xFF == ord('q'):
                break
        else:
            key = cv.waitKey(-1)
            # Break the infinite loop
            if key & 0xFF == ord('q'):
                break
            # Break the infinite loop
            if key & 0xFF == ord('s'):
                if cv.imwrite(yaml_handle.CURRENT_DIR + '/test/screenshots/{}.jpg'.
                              format(int(time.time())), chessboard):
                    print "Image succesfuly saved!"
                    cv.waitKey(1000)
                else:
                    print "Error: Couldnt save image!"
                continue
            # Cycle through the images
            if key & 0xFF == 81:
                frame_index -= 1
            else:
                frame_index += 1

    # When everything done, release the capture and close all windows
    if USE_LIVE_IMAGES:
        cap.release()
    cv.destroyAllWindows()

if __name__ == '__main__':
    main()
