# Copyright 2018 Julian van Doorn and Kiet van Osnabrugge
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import cv2
import numpy as np
from vector import Vector
from pylab import array, plot, show, axis, arange, figure, uint8 
import util

## Calibration class with routines to calibrate the camera
# @details
# Used for calibrating the vision algorithm. The calibration can
# be stored in memory independent of a Vision instance. It also
# encourages the use of OOP.
class Calibration:
    ## Instance constructor
    # @details
    # Calibration instance constructor. Takes
    # a VideoCapture to use as source. Initializes
    # all used member attributes as None.
    # @param videoCapture OpenCV's VideoCapture to use as frame source
    def __init__(self, videoCapture):
        self.videoCapture = videoCapture
        self.redMotor = None
        self.greenMotor = None
        self.blueMotor = None
        self.center = None

    ## Finds the servo stickers on given frame
    # @details
    # Find the servo stickers on the given thresholded
    # frame. This is done by filtering red, green and
    # blue from the image. Then image moments are used
    # to localize the blobs on the image. The returned
    # values may not accurately represent the location
    # of the stickers and are only an estimate. It is
    # advised to use hough circles along this function
    # to accurately identify the locations of the stickers 
    # @param thresholdedFrame Input frame to find the stickers on. Must be thresholed per color channel. Else the stickers may not be properly identified.
    def getServoPoints(self, thresholdedFrame):
        thresholdedFrameHSV = cv2.cvtColor(thresholdedFrame, cv2.COLOR_BGR2HSV)

        redMask = cv2.inRange(thresholdedFrameHSV, np.array([-20, 125, 0]), np.array([20, 255, 255]))
        greenMask = cv2.inRange(thresholdedFrameHSV, np.array([50, 0, 0]), np.array([75, 255, 255]))
        blueMask = cv2.inRange(thresholdedFrameHSV, np.array([100, 0, 0]), np.array([200, 255, 255]))

        blueCx = 0
        blueCy = 0
        greenCx = 0
        greenCy = 0
        redCx = 0
        redCy = 0

        M = cv2.moments(redMask)
        if M['m00'] > 0:
            redCx = int(M['m10']/M['m00'])
            redCy = int(M['m01']/M['m00'])

        M = cv2.moments(greenMask)
        if M['m00'] > 0:
            greenCx = int(M['m10']/M['m00'])
            greenCy = int(M['m01']/M['m00'])

        M = cv2.moments(blueMask)
        if M['m00'] > 0:
            blueCx = int(M['m10']/M['m00'])
            blueCy = int(M['m01']/M['m00'])
        
        return Vector(redCx, redCy), Vector(greenCx, greenCy), Vector(blueCx, blueCy)

    ## Generates calibration data for this instance
    # @details
    # Generates the calibration data for this instance.
    # The data is generated by taking a series of frame
    # and calibrate them individually. Then all valid
    # frame calibrations are taken and an average value
    # is calculated. That average value is the final
    # calibration for the Vision instance in this project.
    def calibrateCamera(self):
        red = None
        green = None
        blue = None
        vals = []

        for i in range(100):
            valid, data = self.getCalibrationFrameData()

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

            if valid:
                vals.append(data)

        cv2.destroyWindow("Calibration")

        red = vals[0][0]
        green = vals[0][1]
        blue = vals[0][2]

        for v in vals:
            red += v[0]
            green += v[1]
            blue += v[2]

        red = Vector(int(red.values[0] / len(vals)), int(red.values[1] / len(vals)))
        green = Vector(int(green.values[0] / len(vals)), int(green.values[1] / len(vals)))
        blue = Vector(int(blue.values[0] / len(vals)), int(blue.values[1] / len(vals)))

        self.redMotor = red
        self.greenMotor = green
        self.blueMotor = blue
        self.center = self.redMotor + self.greenMotor + self.blueMotor
        self.center = Vector(int(self.center.values[0] / 3), int(self.center.values[1] / 3))
    
    ## Generates a calibration for one frame
    # @details
    # Takes a frame from the provided VideoCapture in this 
    # class's constructor. Then applies some thresholding
    # and tries to identify the servo stickers on the frame.
    # A status boolean is returned whether the servo stickers
    # are estimated to be correctly identified alongside the
    # sticker locations. Opens an OpenCV window named "Calibration"
    # you can use cv2.destroyWindow("Calibration") to close
    # this window once calibration was completed.
    # @return True Servos are likely correctly detected
    # @return False Servos are likely not correctly detected
    # @return (RedServo, GreenServo, BlueServo) Tuple containing the locations of the identified servos
    def getCalibrationFrameData(self):
        ret, frame = self.videoCapture.read()

        thresholdedFrame = util.preprocessImage(frame)
        ret, thresholdedFrame =  cv2.threshold(thresholdedFrame,64,255,cv2.THRESH_BINARY)

        red, green, blue = self.getServoPoints(thresholdedFrame)

        thresholdedFrameGray = cv2.cvtColor(thresholdedFrame, cv2.COLOR_RGB2GRAY)
        ret, thresholdedFrameGray = cv2.threshold(thresholdedFrameGray,200,255,cv2.THRESH_BINARY)

        circles = cv2.HoughCircles(thresholdedFrameGray, cv2.HOUGH_GRADIENT,4,50,
                                    param1=50,param2=30,minRadius=0,maxRadius=30)

        if circles is not None and len(circles[0]) >= 3:
            red = Calibration.bindHoughCirclesToColorEstimation(circles, red)
            green = Calibration.bindHoughCirclesToColorEstimation(circles, green)
            blue = Calibration.bindHoughCirclesToColorEstimation(circles, blue)

            cv2.circle(thresholdedFrame, red.values, 4, (0,0,0), -1)
            cv2.circle(thresholdedFrame, green.values, 4, (0,0,0), -1)
            cv2.circle(thresholdedFrame, blue.values, 4, (0,0,0), -1)

            middlePoint = red + green + blue
            middlePoint = Vector(int(middlePoint.values[0] / 3), int(middlePoint.values[1] / 3))

            cv2.circle(thresholdedFrame, middlePoint.values, 4, (0,0,255), -1)

            cv2.line(thresholdedFrame, red.values, middlePoint.values, (0,0,255), 1, cv2.LINE_AA)
            cv2.line(thresholdedFrame, green.values, middlePoint.values, (0,255,0), 1, cv2.LINE_AA)
            cv2.line(thresholdedFrame, blue.values, middlePoint.values, (255,0,0), 1, cv2.LINE_AA)

            if circles is not None:
                circles = np.uint16(np.around(circles))
                for i in circles[0,:]:
                    cv2.circle(thresholdedFrame,(i[0],i[1]),i[2],(0,255,0),2)
                    cv2.circle(thresholdedFrame,(i[0],i[1]),2,(0,0,0),3)

            cv2.imshow("Calibration", thresholdedFrame)

            calibration = (red, green, blue)

            return len(circles[0]) >= 3, calibration
        else:
            return False, (Vector(0, 0), Vector(0, 0), Vector(0, 0))

    ## Binds the closest hough circle to the estimated position
    # @details
    # Binds the closest hough circle to the estimated position.
    # This is useful for binding an id to a hough circle. Since
    # a hough circle does not provide any context about the circle,
    # only where the circle is. So this function is used for tying
    # an exact location to an estimate identification.
    # @param houghCircles Array returned by cv2.HoughCircles(...)
    # @param estimatedPosition Estimated position of a hough circle
    @staticmethod
    def bindHoughCirclesToColorEstimation(houghCircles, estimatedPosition):
        closestHoughCircle = None

        if houghCircles is not None:
            houghCircles = np.uint16(np.around(houghCircles))

            for circle in houghCircles[0,:]:
                coords = Vector(circle[0], circle[1])
                rad = circle[2]

                if closestHoughCircle is None:
                    closestHoughCircle = (coords, rad)
                else:
                    if (estimatedPosition - coords).norm() < (estimatedPosition - closestHoughCircle[0]).norm():
                        closestHoughCircle = (coords, rad)

        if closestHoughCircle is not None:
            return closestHoughCircle[0]
        else:
            return estimatedPosition