#!/usr/bin/env python
"""Script containing function for camera calibration and image proccessing"""
import glob
import math
import time
import copy
import rospy
import numpy as np
import cv2 as cv

import include.yaml_data_handling as yaml_handle

# Debug mode
DEBUG = False

# Get Camera ID from ROS Param Server
CAMERA_ID = rospy.get_param("robocheckers/camera_id")


def show_raw_image(draw_middle=False):
    """Only shows raw camera image for visual feedback during lense and aperture adjustion"""
    # Camera selection
    cap = cv.VideoCapture(CAMERA_ID)

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

        # Return 0 if camera not found
        if not ret:
            print "Error: Camera not found!"
            rospy.sleep(5)
            break

        if draw_middle:
            x = frame.shape[1]
            y = frame.shape[0]
            cv.circle(frame, (x/2, y/2), radius=5, color=(0, 0, 255), thickness=-1)

        cv.imshow('Frame', frame)

        if cv.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv.destroyAllWindows()

    return


def save_live_images(num):
    """
    Snaps and saves number of specified images
    :param num: Number of imagase to be saved
    :return: True/False signalizing succes of the process
    """
    # If number of images to save is 0 or less, return False
    if num <= 0:
        return False

    # Camera selection
    cap = cv.VideoCapture(CAMERA_ID)

    print "Press `s` to save image or `q` to exit."

    while num > 0:

        # Capture frame-by-frame
        ret, frame = cap.read()

        # Return 0 if camera not found
        if not ret:
            print "Error: Camera not found!"
            rospy.sleep(5)
            cap.release()
            return False

        # Show image
        cv.imshow('Frame', frame)
        key = cv.waitKey(1) & 0xFF

        # Press for keyboard imput to save the image or end break
        if key:
            if key == ord('s'):
                cv.imwrite('../calibration/images/{}.jpg'.format(int(time.time())), frame)
                print "Image saved!\nPress `s` to save next image or `q` to exit."

                # Reduce the num value
                num -= 1

            if key == ord('q'):
                break

    cap.release()
    cv.destroyAllWindows()
    return True


def get_camera_params():
    """
    Gets camera calibration data from stored images
    :return: True/False signalizing succes of the process and calibration data
    """
    # Termination criteria
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # Prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    objp = np.zeros((7*7, 3), np.float32)
    objp[:, :2] = np.mgrid[0:7, 0:7].T.reshape(-1, 2)

    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.

    images = glob.glob(yaml_handle.CURRENT_DIR + '/calibration/images/*.jpg')

    found = 0
    # List all loaded images and find chessboard corners
    for fname in images:
        img = cv.imread(fname)
        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

        # Find the chess board corners
        ret, corners = cv.findChessboardCorners(gray, (7, 7), None)

        # If found, add object points, image points (after refining them)
        if ret:
            objpoints.append(objp)   # Certainly, every loop objp is the same, in 3D.
            corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            imgpoints.append(corners2)

            # Draw and display the corners
            img = cv.drawChessboardCorners(img, (7, 7), corners2, ret)
            found += 1
            if DEBUG:
                print "Valid img:", fname
                cv.imshow('img', img)
                cv.waitKey(500)
        else:
            if DEBUG:
                print "Invalid img:", fname

    print "Number of images used for calibration:", found

    # When everything done, close all windows
    cv.destroyAllWindows()

    # calibration
    ret, mtx, dist, _, _ = cv.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

    if ret:
        return True, mtx, dist
    return False, None, None


def calibrate_frame(frame, mtx, dist):
    """Calibrates given frame by saved calibration data from yaml file and returns it"""
    height, width = frame.shape[:2]
    newcameramtx, _ = cv.getOptimalNewCameraMatrix(mtx, dist, (width, height), 1, (width, height))

    # Calculate center of the frame and prepare it for undostortPoints function
    center = (width/2-200, height/2)
    #center = np.array([(width/2, height/2)])
    #center = np.float32(center[:, np.newaxis, :])

    # Undistort frame and center
    #dst = cv.undistort(frame, mtx, dist, None)
    #dst = cv.undistort(frame, mtx, dist, None, newcameramtx)
    mapx, mapy = cv.initUndistortRectifyMap(mtx, dist, None, newcameramtx, (width, height), 5)
    dst = cv.remap(frame, mapx, mapy, cv.INTER_LINEAR)
    #center = cv.undistortPoints(np.array(center), mtx, dist, None, newcameramtx)

    # crop the image and move the center of chessboard to the right position
    #x, y, w, h = roi
    #dst = dst[y:y+h, x:x+w]
    dst = dst[:, 200:1080]
    #center[0, 0, 0] -= 200

    if DEBUG:
        cv.imshow("Not calibrated", frame)
        cv.waitKey(100)
        cv.imshow("Calibrated", dst)
        cv.waitKey(100)

        raw_input("Press `Enter` to close windows")

        cv.destroyAllWindows()

    # Return calibrated image
    #return dst, tuple(center[0, 0])
    return dst, center


def get_corners(frame):
    """Finds chessboard corners in frame"""
    # Termination criteria
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # Convert to gray
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)

    # Find the chessboard corners in image and draw them
    ret, corners = cv.findChessboardCorners(gray, (7, 7), flags=cv.CALIB_CB_ADAPTIVE_THRESH)
    if ret:
        corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        cv.drawChessboardCorners(frame, (7, 7), corners2, ret)

    if DEBUG:
        cv.imshow("Corners", frame)
        cv.waitKey(100)

        raw_input("Press `Enter` to close windows")

        cv.destroyAllWindows()

    # After successful save, break the main cycle
    if ret:
        print "Chessboard fixed corners found successfuly!"
        return corners2
    print "Error: Chessboard fixed corners not found!"
    return None


def get_chessboard_data(corners):
    """Gets chessboard position in real world and chessboard dimensions data"""
    # Separate x and y coordinations
    corners = corners[:, 0]
    corners_x = corners[:, 1]
    corners_y = corners[:, 0]

    # Predefine distance variable
    distance = []

    # Find 1st random corner point
    for i in corners:
        tmp = np.sqrt((corners[0][0]-i[0]) ** 2 + (corners[0][1]-i[1]) ** 2)
        distance.append(tmp)

    # Save 1st corner point into the extremes list
    extremes = []
    extremes.append(corners[distance.index(max(distance))])

    # Clear the distance variable
    distance = []

    # Find 2st corner point opposite to the 1st one
    for i in corners:
        tmp = np.sqrt((extremes[0][0]-i[0]) ** 2 + (extremes[0][1]-i[1]) ** 2)
        distance.append(tmp)

    # Add the 2nd corner point into the extremes list
    extremes.append(corners[distance.index(max(distance))])

    # Find center of the square
    center = ((extremes[0][0] + extremes[1][0]) / 2, (extremes[0][1] + extremes[1][1]) / 2)
    print center

    # Clear the distance variable
    distance = []

    # Calculate all distances from the center
    for i in range(49):
        tmp = np.sqrt((center[0]-corners[i][0]) ** 2 + (center[1]-corners[i][1]) ** 2)
        distance.append((tmp, i))

    # Sort by distance in descending order
    distance = sorted(distance, reverse=True, key=lambda x: x[0])

    # Clear extremes
    extremes = []

    # Save four the most distant points from the center
    for i in range(4):
        pom = corners[distance[i][1]]
        extremes.append(pom.tolist())

    if DEBUG:
        print "Extremes:\n", extremes

    # Sort x and y coordinations
    corners_x = sorted(corners_x, reverse=True)
    corners_y = sorted(corners_y, reverse=True)

    # Prepare arrays for average positions for every line of corners
    avg_x = []
    avg_y = []

    # Calculate average x and y positions for every line of corners imaginable
    for i in range(0, 49, 7):
        avg_x.append(sum(corners_x[i:i+7])/7)
        avg_y.append(sum(corners_y[i:i+7])/7)

    # Prepare arrays for average distances between corners of the chessboard
    avg_dx = []
    avg_dy = []

    # Calculate average distances between adjecent corners
    for i in range(6):
        avg_dx.append(avg_x[i] - avg_x[i+1])
        avg_dy.append(avg_y[i] - avg_y[i+1])

    # Discard unnecessary
    avg_dx = avg_dx[:6]
    avg_dy = avg_dy[:6]

    # Calculate average distance from average distances between adjecent corners
    avg_dx = sum(avg_dx)/6
    avg_dy = sum(avg_dy)/6
    avg_tile = (avg_dx + avg_dy)/2

    #print corners_x
    #print corners_y
    #print "Corners:\n", corners

    print avg_dx
    print avg_dy
    print avg_tile

    return extremes, avg_dx, avg_dy


def get_center(corners):
    """Calculate center from 4 points"""
    # Calculate the center of the 4 points
    center = ((corners[0][0] + corners[1][0] + corners[2][0] + corners[3][0])/4,
              (corners[0][1] + corners[1][1] + corners[2][1] + corners[3][1])/4)

    return center


def get_angle(corners):
    """Calculate apexes of the chessboard and rotation"""
    # Sort corners
    corners = sorted(corners, key=lambda x: x[0])

    # Calculate angle of the chessboard in radians
    angle = math.atan(abs((corners[0][0] - corners[1][0])/(corners[0][1] - corners[1][1])))

    if DEBUG:
        print "Angle:", angle, "rad"

    return angle


def adjust_corners(corners, center, angle, dx, dy):
    """Rotates the corners and calculate the apexes"""
    # Convert angle from radians to degrees
    angle = -angle

    # For some ununderstandable reason, I am unable to COPY the list corners,
    # cause Python 2 and I am frustrated enough to just do it like this.
    # It still doesnt work, so whatever
    apexes = sorted(corners, key=lambda x: x[1])

    # Rotate the corners and calculate the apexes of the chessboard
    for i in range(4):
        # Calculate the coordinates of the corners from the center
        apexes[i][0] -= center[0]
        apexes[i][1] -= center[1]

        # Rotating the corners around the center
        apexes[i][0] = apexes[i][0] * math.cos(angle) - apexes[i][1] * math.sin(angle)
        apexes[i][1] = apexes[i][0] * math.sin(angle) + apexes[i][1] * math.cos(angle)

        # Calculate the apexes of the chessboard
        if abs(apexes[i][0] + dx) > abs(apexes[i][0] - dx):
            apexes[i][0] += dx
        else:
            apexes[i][0] -= dx

        if abs(apexes[i][1] + dy) > abs(apexes[i][1] - dy):
            apexes[i][1] += dy
        else:
            apexes[i][1] -= dy

        # Calculate the original 0, 0 coordination position
        apexes[i][0] += center[0]
        apexes[i][1] += center[1]

    return apexes


def rotate_frame(frame, chessboard_center, angle, camera_center=None):
    """Rotates the frame and camera_center"""
    # Rotating the frame
    rot_mat = cv.getRotationMatrix2D(chessboard_center, math.degrees(angle), 1.0)
    result = cv.warpAffine(frame, rot_mat, frame.shape[1::-1], flags=cv.INTER_LINEAR)

    if camera_center is None:
        return result
        
    # Convert tuple to list (cause tuple sucks)
    tmp = list(camera_center)

    # Calculate the coordinates of the corners from the center
    tmp[0] -= chessboard_center[0]
    tmp[1] -= chessboard_center[1]

    # Rotating the camera center
    tmp[0] = int(round(tmp[0] * math.cos(angle) - tmp[1] * math.sin(angle)))
    tmp[1] = int(round(tmp[0] * math.sin(angle) + tmp[1] * math.cos(angle)))

    # Calculate the original 0, 0 coordination position
    tmp[0] += chessboard_center[0]
    tmp[1] += chessboard_center[1]

    # Change back to tuple as opencv guides
    camera_center = tuple(tmp)

    return result, camera_center


def adjust_frame(frame, apexes, camera_center=None, offset=None):
    """Finds the chessboard and crops the frame"""
    apexes = copy.deepcopy(apexes)
    for i in range(4):
        apexes[i][0] = int(round(apexes[i][0]))
        apexes[i][1] = int(round(apexes[i][1]))

    # Prepare the maximum and minimum variables
    min_x = frame.shape[0]
    max_x = -1
    min_y = frame.shape[1]
    max_y = -1

    # Get the maximums of the chessboard corners
    for item in apexes:
        if item[0] < min_x:
            min_x = item[0]
        if item[0] > max_x:
            max_x = item[0]
        if item[1] < min_y:
            min_y = item[1]
        if item[1] > max_y:
            max_y = item[1]

    # Expand limits by offset
    if offset:
        min_x -= offset
        max_x += offset
        min_y -= offset
        max_y += offset

    if DEBUG:
        print "Apexes:", apexes

    #result = frame[min_x:max_x, min_y:max_y]

    # Define the polygon coordinates to use or the crop
    #apexes[2], apexes[3] = apexes[3], apexes[2]

    croped_image = frame
    # Go over the points in the image if thay are out side of the emclosing rectangle put zero
    # if not check if thay are inside the polygon or not
    #croped_image = np.zeros_like(frame)
    #for y in range(min_x, max_x):
    #    for x in range(min_y, max_y):
    #        if cv.pointPolygonTest(np.array(apexes, np.int32), (y, x), False) >= 0:
    #            croped_image[x, y] = frame[x, y]

    # If debug, draw apexes and center
    if DEBUG:
        for item in apexes:
            item = np.array(item)
            item = tuple(item.astype(int))
            croped_image = cv.circle(croped_image, item, radius=5, color=(0, 0, 255), thickness=-1)

    # Now we can crop again just the envloping rectangle
    result = croped_image[min_y:max_y, min_x:max_x]

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

    # Calculating new camera center
    camera_center = list(camera_center)
    camera_center[0] -= min_x
    camera_center[1] -= min_y
    camera_center = tuple(camera_center)

    return result, camera_center


def main():
    """Main function"""
    # Show raw image
    if DEBUG:
        show_raw_image(True)

    #yaml_handle.load_dimensions_data()
    #yaml_handle.save_dimensions_data(55.2, 55.8, 447)

    # Save images
    #save_live_images(4)
    #return

    #ret, mtx, dist = get_camera_params()
    #yaml_handle.save_camera_params(mtx, dist)

    mtx, dist = yaml_handle.load_camera_params()

    cap = cv.VideoCapture(CAMERA_ID)

    corners = None

    while corners is None:
        ret, frame = cap.read()

        # Return 0 if camera not found
        if not ret:
            print "Error: Camera not found!"
            rospy.sleep(5)
            cap.release()
            return False

        dst, _ = calibrate_frame(frame, mtx, dist)

        #cv.imwrite('../calibration/images/test.jpg', dst)
        #dst = cv.imread('../calibration/images/test.jpg')

        corners = get_corners(dst)

    four_corners, dx, dy = get_chessboard_data(corners)

    center = get_center(four_corners)

    angle = get_angle(four_corners)

    center = np.array(center)
    center = tuple(center.astype(int))

    # Find the apexes of the keaboard and rotate them
    apexes = adjust_corners(four_corners, center, angle, dx, dy)

    # Rewrite old data with new
    yaml_handle.save_chessboard_params(center, angle, apexes)
    yaml_handle.center, angle, apexes = yaml_handle.load_chessboard_params()

    dst = rotate_frame(dst, center, angle)
    dst = adjust_frame(dst, apexes)

    # Calculating chessboard lenghts in pixels
    apexes = sorted(apexes, reverse=True, key=lambda x: x[0])
    chessboard_x = abs((apexes[0][0] + apexes[1][0])/2 - (apexes[2][0] +  apexes[3][0])/2)
    apexes = sorted(apexes, reverse=True, key=lambda x: x[1])
    chessboard_y = abs((apexes[0][1] + apexes[1][1])/2 - (apexes[2][1] +  apexes[3][1])/2)
    # Calculating average chessboard len in pixels
    chessboard_px = (chessboard_x + chessboard_y)/2

    print chessboard_x, chessboard_y

    # Save dimensions data
    yaml_handle.save_dimensions_data(dx, dy, chessboard_px)
    yaml_handle.load_dimensions_data()

    if DEBUG:
        cv.imshow('Center', dst)
        cv.waitKey(1)
        print "Press any key to exit."
        cv.waitKey(-1)

    # When function terminated, release the capture and values
    cap.release()
    cv.destroyAllWindows()


if __name__ == '__main__':
    DEBUG = True
    main()
