#Universal Vision Processor For FRC Team 1391
#Designed for use with Raspberry Pi 3 with OpenCV
import sys
import os
import cv2
import numpy
from enum import Enum
import math
from networktables import NetworkTables
import time
import threading
import copy
import logging
import datetime
import socket

############################################################################################
#CHANGE VARIABLES HERE AND ONLY HERE#

#define all hsl values here.
hl = 45
hh = 70
sl = 150
sh = 255
ll = 92
lh = 173

#define roboRIO IP address here
rioIP = '10.13.91.2'
TCPPort = 5011
TCPSocketTimeout = 1

#define all camera IPs and addresses here (as strings). Note: Camera address and IP might be different; in this case using an AXIS M1011, the camera address specifically pulls a mjpeg stream.
gCamIP = '10.13.91.3'
gCamAddr = 'http://10.13.91.3/axis-cgi/mjpg/video.cgi?.mjpg'

bCamIP = '10.13.91.3'
bCamAddr = 'http://10.13.91.3/axis-cgi/mjpg/video.cgi?.mjpg'

#define camera attributes (including placement, tech specs, etc.) here
bCamAngle = 0
bCamElevation = 0

bCamXFOV = 47
bCamYFOV = 35.25
bCamXRes = 640
bCamYRes = 480

gCamAngle = 0
gCamElevation = 0

gCamXFOV = 47
gCamYFOV = 35.25
gCamXRes = 640
gCamYRes = 480

bTarget1Height = 4
bTarget1Width = 15
bTarget1Elevation = 84

bTarget2Height = 2
bTarget2Width = 15
bTarget2Elevation = 80

gTargetHeight = 5
gTargetWidth = 2
gTargetElevation = 10.75

######################################################################################33
class GripPipeline:
    """
    An OpenCV pipeline generated by GRIP.
    """

    def __init__(self, hl, hh, sl, sh, ll, lh):
        """initializes all values to presets or None if need to be set
        """

        self.__hsl_threshold_hue = [hl, hh]
        self.__hsl_threshold_saturation = [sl, sh]
        self.__hsl_threshold_luminance = [ll, lh]

        self.hsl_threshold_output = None

        self.__find_contours_input = self.hsl_threshold_output
        self.__find_contours_external_only = True

        self.find_contours_output = None

        self.__filter_contours_contours = self.find_contours_output
        self.__filter_contours_min_area = 75.0
        self.__filter_contours_min_perimeter = 0
        self.__filter_contours_min_width = 0
        self.__filter_contours_max_width = 1000.0
        self.__filter_contours_min_height = 0
        self.__filter_contours_max_height = 1000
        self.__filter_contours_solidity = [0.0, 100]
        self.__filter_contours_max_vertices = 1000000
        self.__filter_contours_min_vertices = 0
        self.__filter_contours_min_ratio = 0
        self.__filter_contours_max_ratio = 1000

        self.filter_contours_output = None

        self.good_contours = None
        self.sorted_contours = None

    def process(self, source0):
        """
        Runs the pipeline and sets all outputs to new values.
        """
        # Step HSL_Threshold0:
        self.__hsl_threshold_input = source0
        (self.hsl_threshold_output) = self.__hsl_threshold(self.__hsl_threshold_input, self.__hsl_threshold_hue, self.__hsl_threshold_saturation, self.__hsl_threshold_luminance)

        # Step Find_Contours0:
        self.__find_contours_input = self.hsl_threshold_output
        (self.find_contours_output) = self.__find_contours(self.__find_contours_input, self.__find_contours_external_only)

        # Step Filter_Contours0:
        self.__filter_contours_contours = self.find_contours_output
        self.good_contours = self.__filter_contours(self.__filter_contours_contours, self.__filter_contours_min_area, self.__filter_contours_min_perimeter, self.__filter_contours_min_width, self.__filter_contours_max_width, self.__filter_contours_min_height, self.__filter_contours_max_height, self.__filter_contours_solidity, self.__filter_contours_max_vertices, self.__filter_contours_min_vertices, self.__filter_contours_min_ratio, self.__filter_contours_max_ratio)

        self.sorted_contours = sorted(self.good_contours, key=cv2.contourArea, reverse=True)

        if len(self.sorted_contours) >= 2:
            contour1Atrib = self.contourReport(self.sorted_contours[0])
            contour2Atrib = self.contourReport(self.sorted_contours[1])
            return contour1Atrib, contour2Atrib, (self.hsl_threshold_output)
        else:
            return False, (self.hsl_threshold_output)

    @staticmethod
    def __hsl_threshold(input, hue, sat, lum):
        """Segment an image based on hue, saturation, and luminance ranges.
        Args:
            input: A BGR numpy.ndarray.
            hue: A list of two numbers the are the min and max hue.
            sat: A list of two numbers the are the min and max saturation.
            lum: A list of two numbers the are the min and max luminance.
        Returns:
            A black and white numpy.ndarray.
        """
        #cv2.imshow('opencv input', input)
        #print 'hsl'
        out = cv2.cvtColor(input, cv2.COLOR_BGR2HLS)
        return cv2.inRange(out, (hue[0], lum[0], sat[0]),  (hue[1], lum[1], sat[1]))

    @staticmethod
    def __find_contours(input, external_only):
        """Sets the values of pixels in a binary image to their distance to the nearest black pixel.
        Args:
            input: A numpy.ndarray.
            external_only: A boolean. If true only external contours are found.
        Return:
            A list of numpy.ndarray where each one represents a contour.
        """
        #print 'coutouring'
        if(external_only):
            mode = cv2.RETR_EXTERNAL
        else:
            mode = cv2.RETR_LIST
        method = cv2.CHAIN_APPROX_SIMPLE
        im2, contours, hierarchy =cv2.findContours(input, mode=mode, method=method)
        #cv2.imshow('hsl', input)
        #print contours
        return contours

    @staticmethod
    def __filter_contours(input_contours, min_area, min_perimeter, min_width, max_width,
                        min_height, max_height, solidity, max_vertex_count, min_vertex_count,
                        min_ratio, max_ratio):
        """Filters out contours that do not meet certain criteria.
        Args:
            input_contours: Contours as a list of numpy.ndarray.
            min_area: The minimum area of a contour that will be kept.
            min_perimeter: The minimum perimeter of a contour that will be kept.
            min_width: Minimum width of a contour.
            max_width: MaxWidth maximum width.
            min_height: Minimum height.
            max_height: Maximimum height.
            solidity: The minimum and maximum solidity of a contour.
            min_vertex_count: Minimum vertex Count of the contours.
            max_vertex_count: Maximum vertex Count.
            min_ratio: Minimum ratio of width to height.
            max_ratio: Maximum ratio of width to height.
        Returns:
            Contours as a list of numpy.ndarray.
        """
        #print 'filtering'
        output = []
        for contour in input_contours:
            #print 'contour?'
            x,y,w,h = cv2.boundingRect(contour)
            if (w < min_width or w > max_width):
                continue
            if (h < min_height or h > max_height):
                continue
            area = cv2.contourArea(contour)
            if (area < min_area):
                continue
            if (cv2.arcLength(contour, True) < min_perimeter):
                continue
            hull = cv2.convexHull(contour)
            solid = 100 * area / cv2.contourArea(hull)
            if (solid < solidity[0] or solid > solidity[1]):
                continue
            if (len(contour) < min_vertex_count or len(contour) > max_vertex_count):
                continue
            ratio = (float)(w) / h
            if (ratio < min_ratio or ratio > max_ratio):
                continue
            #print contour
            output.append(contour)
            #print 'good contour'
        #print output
        return output

def contourReport(contour):
    x, y, w, h = cv2.boundingRect(contour)
    centerx = x + (w/2)
    centery = y + (h/2)
    aspect = float(w)/h
    contourAtrib = {'x':x, 'y':y, 'w':w, 'h':h, 'centerx':centerx, 'centery':centery, 'aspect': aspect}
    return contourAtrib

class Target:

    def __init__(self, width, height, elevation):
        self.width = width
        self.height = height
        self.elevation = elevation

    #returns returns horizontal angle offset between a given point in a frame and the center of the camer lens
    def getAngleX(camera, rectCenterX):
        return camera.xFOV*(float(rectCenterX - camera.halfXRes)/camera.halfXRes)

    #returns vertical angle offset between a given point in a frame and the center of the camer lens
    def getAngleY(camera, rectCenterY):
        return camera.yFOV*(float(rectCenterY - camera.halfYRes)/camera.halfYRes)

    #returns distance to target based on target height
    def getDistHeight(self, camera, rectHeight):
        return float(self.height*camera.yRes)/(2*rectHeight*(math.tan(math.radians(camera.halfYFOV))))

    #Returns distance to target based on target width
    def getDistWidth(self, camera, rectWidth):
        return float(self.width*camera.xRes)/(2*rectWidth*(math.tan(math.radians(camera.halfXFOV))))

    #returns distance to target
    def getDistElevation(self, camera, rectLowerY):
        return float(self.elevation)/math.tan(math.radians(self.getAngleY(rectLowerY)-camera.angle))

class Camera:

    #returns camera object with given perameters
    def __init__(self, ip, addr, angle, elevation, xRes, yRes, xFOV, yFOV):
        self.xFOV = xFOV
        self.yFOV = yFOV
        self.xRes = xRes
        self.yRes = yRes
        self.halfXFOV = float(xFOV) / 2
        self.halfYFOV = float(yFOV) / 2
        self.halfXRes = xRes / 2
        self.halfYRes = yRes / 2
        self.ip = ip
        self.addr = addr
        self.angle = angle
        self.elevation = elevation
        self.cv2CameraObject = False
        self.ip = ip
        self.isConnected = False

    #Passed camera object, attempts to connect x times
    def connect(self, tries):
        global logger

        for i in range(tries):
            logger.info('Connecting to camera...')
            try:
                self.cv2CameraObject = cv2.VideoCapture(self.addr)
            except:
                logger.warning('Camera inaccessible! %s tries left.', str(range(tries)-i))
                continue

            camTestBool, val = self.cv2CameraObject.read()

            if camTestBool:
                self.isConnected = True
                logger.info('Camera connected.')
                break
            else:
                logger.warning('Camera inaccessible! %s tries left.', str(range(tries)-i))
                self.isConnected = False
                self.cv2CameraObject = False

    #Attempts to return a frame from given object
    def getFrame(self):
        global logger

        logger.debug('Retrieving frame from camera...')
        try:
            retVal, img = self.cv2CameraObject.read()
        except:
            logger.warning('Frame unavailable!')
            self.isConnected = False
            return False
        if retVal:
            logger.debug('Frame retrieved.')
            return img
        else:
            logger.warning('Frame unavailable!')
            self.isConnected = False
            return False

        #determines whether connection to a network device (roboRIO, camera) is possible
    def ping(self):
        if os.system("ping " + str(self.ip) + " -c 1 -W 3") == 0:
            self.isConnected = True
            return True
        else:
            self.isConnected = False
            return False

    #returns result of last ping
    def getConnectionStatus(self):
        return self.isConnected

#used to communicate between program threads intead of variables.
#any variables that need to be accessed by both manageCameras() and manageComputation() should instead be instances of MultiThreadVariable; this object and its associated read() and write() functions have implemented memory protection
class MultiThreadVariable:
    global logger

    def __init__(self, value):
        self.value = value
        self.Lock = threading.Lock()

    def write(self, value):
        logger.debug('Writing to MTV')
        with self.Lock:
            self.value = copy.copy(value)

    def read(self):
        logger.debug('Reading from MTV')
        with self.Lock:
            return self.value

def initialize():
    #use this method to connect to networktables, instantiate and connect to cameras, create an instance of the pipe process, set debugging and logging perameters, etc; every function that runs once at program bootup should be defined here.
    global isShowFrames
    global isDebugging
    global isNetworked
    global isWarning
    global logger

    isShowFrames = False
    isDebugging = False
    isNetworked = True
    isWarning = False

    if len(sys.argv) > 1:
        for arg in sys.argv:
            if arg == "-s":
                isShowFrames = True
            if arg == "-d":
                isDebugging = True
            if arg == '-n':
                isNetworked = False
            if arg == '-w':
                isWarning = True
    FORMAT = logging.Formatter(fmt='%(levelname)s : %(threadName)s @ %(asctime)s : %(message)s')
    FILE_HANDLER = logging.FileHandler('visionlogs/'+(str(datetime.datetime.now().strftime('%m-%d-%y@%H:%M:%S'))+'.log'))
    FILE_HANDLER.setFormatter(FORMAT)
    STREAM_HANDLER = logging.StreamHandler(stream=sys.stdout)
    STREAM_HANDLER.setFormatter(FORMAT)
    logger = logging.getLogger()
    logger.addHandler(FILE_HANDLER)
    logger.setLevel(logging.INFO)
    if isDebugging:
        logger.addHandler(STREAM_HANDLER)
        logger.setLevel(logging.DEBUG)
    if isWarning:
        logger.setLevel(logging.WARNING)
    logger.info('Debugging perameters set. Logging system initialized.')

    bTarget1 = Target(bTarget1Width, bTarget1Height, bTarget1Elevation)
    bTarget2 = Target(bTarget2Width, bTarget2Height, bTarget2Elevation)
    gTarget = Target(gTargetWidth, gTargetHeight, gTargetElevation)

    gCamera = Camera(gCamIP, gCamAddr, gCamAngle, gCamElevation, gCamXRes, gCamYRes, gCamXFOV, gCamYFOV)
    bCamera = Camera(bCamIP, bCamAddr, bCamAngle, bCamElevation, bCamXRes, bCamYRes, bCamXFOV, bCamYFOV)

    logger.info('Pinging camera...')

    while gCamera.ping() == False and bCamera.ping == False:
        if gCamera.ping():
            gCamera.connect(5)
            break
        elif bCamera.ping():
            bCamera.connect(5)
            break
        time.sleep(1)
        logger.warning('Camera inaccessible!')



    pipe = GripPipeline()

    #creates all MultiThreadVariables for threasafe communication
    sharedImageTuple = MultiThreadVariable((False, False))
    sharedDist = MultiThreadVariable(0)
    sharedAngle = MultiThreadVariable(0)
    sharedIsTargetInFrame = MultiThreadVariable(0)
    sharedVisionState = MultiThreadVariable(False)
    #As a final step in this method, create the two program threads and initialize both.

    logger.info('Initializing camera and computation threads.')
    cameraThread = Thread.threading(target=manageCamera, args=(gCamera, bCamera, sharedImageTuple), name='Camera Manager')
    computeThread = Thread.threading(target=manageComputation, args=(pipe, gTarget, bTarget1, bTarget2, gCamera, bCamera, sharedImageTuple, sharedDist, sharedAngle, sharedIsTargetInFrame, sharedVisionState), name='Processing Manager')
    TCPThread = Thread.threading(target=manageTCP, args=(rioIP, TCPPort, TCPSocketTimeout, sharedDist, sharedAngle, sharedIsTargetInFrame), name='TCP Manager')
    cameraThread.start()
    computeThread.start()
    TCPThread.start()

def manageCamera(gCamera, bCamera, sharedImageTuple, sharedVisionState):
    global logger
    #this loop will run continuously at the end of the initialize() sequence.
    #It should contain all code pertaining to pulling frames from a camera. It should also include all code that tries to reconnect to disconnected cameras, and pass the pulled frame as a MultiThreadVariable to the computationThread
    logger.info('Camera Thread initialized.')
    while True:
        if sharedVisionState.read():
            img = bCamera.getFrame()
            if img == False:
                logger.warning("Image unretrievable from boiler camera; reconnecting.")
                bCamera.connect(1)
                else:
                    sharedImageTuple.write((img, True))

        else:
            img = gCamera.getFrame()
            if img == False:
                logger.warning("Image unretrievable gear camera; reconnecting.")
                gCamera.connect(1)
                else:
                    sharedImageTuple.write((img, True))



def manageComputation(pipe, gTarget, bTarget1, bTarget2, gCamera, bCamera, sharedImageTuple, sharedDist, sharedAngle, sharedIsTargetInFrame, sharedVisionState):
    #this loop will run continuously at the end of the initialize() sequence.
    #It should contain all code that must be run to process each frame.
    global isShowFrames
    global isDebugging

    if isShowFrames:
        cv2.namedWindow("image")
        cv2.namedWindow('hsl')
    logger.info('Computation Thread initialized.')
    while True:
        img, isNewFrame = sharedImageTuple.read()
        sharedImageTuple.read(None, False)
        if isNewFrame == False:
            logger.debug('No new frame.')
            continue
        logger.info('Processing frame...')
        logger.debug('Finding contours...')
        outputs, hsl = pipe.process(img)
        contours = sorted(outputs, key=cv2.contourArea, reverse=True)
        logger.debug('Contours processed')

        if len(contours) >= 1:
            logger.debug('Processing contours...')
            contourAtrib1 = contourReport(contours[0])
            contourAtrib2 = contourReport(contours[1])
            if sharedVisionState.read():
                angle = bTarget1.getXAngle(bCamera, contourAtrib1[centerx])
                dist = target.getDistElevation(bCamera, contourAtrib1[y])
                isTargetInFrame = True
            else:
                angle = gTarget.getXAngle(gCamera, float(contourAtrib1[centerx] + contourAtrib2[centerx])/2)
                dist = target.getDistHeight(gCamera, contourAtrib[height])
                isTargetInFrame = True
            logger.debug('Contours processed.')
            logger.info('Frame processed.')
            if isShowFrames:
                cv2.rectangle(img, (320, 240), (320, 240), (255, 0, 0), 5)
                cv2.rectangle(img, (contourAtrib[x],contourAtrib[y]), (contourAtrib[x]+contourAtrib[w], contourAtrib[y]+contourAtrib[h]), (255,0,0), 1)
                cv2.rectangle(img, (contourAtrib[centerx], contourAtrib[centery]), (contourAtrib[centerx], contourAtrib[centery]), (0,0,255), 1)

        else:
            logger.info('Not enough contours; no target in frame.')
            angle = 0
            dist = 0
            isTargetInFrame = False

        logger.info('Publishing values to MTV')
        sharedDist.write(dist)
        sharedAngle.write(angle)
        sharedIsTargetInFrame.write(isTargetInFrame)

        if isDebugging:
            print "Dist: " + dist
            print "Angle: " + angle
            print "isTargetInFrame: " + isTargetInFrame

        if isShowFrames:
            cv2.imshow('image', img)
            cv2.imshow('hsl', hsl)

#Communicates with roborio over TCP socket
def manageTCP(addr, port, timeout, sharedDist, sharedAngle, sharedIsTargetInFrame, sharedVisionState):
    while True:
        #creates socekt and connects to roboRIO
        logging.info("Initializing TCP Thread")
        greeter = socket.socket()
        greeter.bind(addr, port)
        logging.info("Listening for connection...")
        greeter.listen(1)

        socket, clientAddr = greeter.accept()
        logging.info("Connection made at %" %clientAddr)

        socket.settimeout(timeout)

        logging.info("Beginning query/response sequence")
        while True:
            try:
                query = socket.recv(1024)
                logging.debug("Query recieved from client")
            except:
                logging.error("Timeout waiting for client query!")
                break
            if query == "q":

                    logging.error("Error sending packet to client!")
            if query == "g":
                sharedVisionState.write(False)
                logging.info("Setting visionState to gear")

            elif query == "b":
                sharedVisionState.write(True)
                logging.info("Setting visionstate to boiler")
            try:
                socket.send(str(sharedDist.read()) + ':' + str(sharedAngle.read()) + ':' + str(sharedIsTargetInFrame.read())+"\n")
                logging.debug("Packet sent to client")
            except:
                break

        logging.warning("Send/recieve exited with errors. Closing sockets and reinitializing.")
        greeter.close()
        socket.close()

initialize()
