import numpy as np
import cv2
from math import sqrt


depth_scale = 0.0010000000474974513
print("Depth Scale is: " , depth_scale)

font = cv2.FONT_HERSHEY_SIMPLEX


color_range =np.array([[20,100,100,25,255,255],[90,250,71,110,255,255],\
                        [80,150,50,90,255,150],[170,190,175,179,255,255],\
                        [12,200,123,19,244,200],[20,00,238,40,255,255],\
                        [0,165,250,8,210,255],[68,120,0,90,255,53]])


active = []
tbc = []

color_c = []
c_h = []

try:
    while True:
        # Get frameset of color and depth
        frames = pipeline.wait_for_frames()
        # Align the depth frame to color frame
        aligned_frames = align.process(frames)
        # Get aligned frames
        aligned_depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()
        # Validate that both frames are valid
        if not aligned_depth_frame or not color_frame:
            continue
        # the depth_image array can be used in order to compute the distance,
        # if the value is multiplied with the the depth_scale value
        depth_image = np.asanyarray(aligned_depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        # color_image = cv2.imread('color.jpeg',1)
        # np.savetxt('distance.txt', depth_image , fmt='%f')
        # cv2.imwrite('color.jpeg' ,color_image)
        gray = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)


        gray = cv2.medianBlur(gray,7)
        img = color_image.copy()
        edges = cv2.Canny(gray,100,25) #TODO parameters to be tuned
        c_h *= 0
        circles = cv2.HoughCircles(gray,cv2.HOUGH_GRADIENT,1,30,
                                    param1=50,param2=25,minRadius=10,maxRadius=50)
        if(circles is not None):
            circles = np.reshape(circles, (np.shape(circles)[1],3))
            for i in range(0,np.shape(circles)[0]):
                cX = circles[i,0]
                cY = circles[i,1]
                r = circles[i,2]
                c_h.append((cX,cY,r))

        color_c *= 0
        hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
        thr = np.zeros((8,720,1280))
        # Quite tedious solution, but each color seems to need its own treatment
        for i in range(0,np.shape(color_range)[0]):
            lower = color_range[i,0:3]
            upper = color_range[i,3:6]
            mask = cv2.inRange(hsv, lower, upper)
            if (i == 3):
                # print('y')
                mask = cv2.dilate(mask, None, iterations=6)
                mask2 = cv2.inRange(hsv,(2,213,60),(13,255,180))
                # mask2 = cv2.dilate(mask2, None, iterations=1)
                mask = mask2 | mask
                mask = cv2.dilate(mask, None, iterations=1)
                mask = cv2.erode(mask, None, iterations=2)
                mask = cv2.dilate(mask, None, iterations=2)

            elif (i==6):
                mask2 = cv2.inRange(hsv,(176,129,230),(179,156,255))
                mask = cv2.dilate(mask, None, iterations=7)
                mask = cv2.erode(mask, None, iterations=2)
                mask = cv2.dilate(mask, None, iterations=3)
                mask = mask2 | mask
                mask = cv2.erode(mask, None, iterations=1)
                mask = cv2.dilate(mask, None, iterations=1)
            elif (i==7):
                mask = cv2.erode(mask, None, iterations=2)
                mask = cv2.dilate(mask, None, iterations=6)
            elif (i==0):
                mask = cv2.erode(mask, None, iterations=1)
                mask = cv2.dilate(mask, None, iterations=4)
            elif (i==1):
                mask = cv2.erode(mask, None, iterations=1)
                mask = cv2.dilate(mask, None, iterations=4)
            elif (i==2):
                mask = cv2.erode(mask, None, iterations=2)
                mask = cv2.dilate(mask, None, iterations=5)
            elif (i==4):
                mask = cv2.erode(mask, None, iterations=0)
                mask = cv2.dilate(mask, None, iterations=7)
            elif (i==5):
                mask = cv2.erode(mask, None, iterations=1)
                mask = cv2.dilate(mask, None, iterations=6)
            thr[i,:,:] = cv2.threshold(mask,150,255,cv2.THRESH_BINARY)[1]
            cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[1]
            # print('start')

            if(len(cnts))> 0:
                for c in cnts:
                    a = cv2.contourArea(c)
                    if (a > 0 and a < 7850):
                        ((x,y), radius) = cv2.minEnclosingCircle(c)
                        M = cv2.moments(c)
                        if M["m00"] != 0:
                            center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
                        else:
                            center = (0, 0)
                        cX = center[0]
                        cY = center[1]
                        r = radius

                        color_c.append((cX,cY,r))

        c_final = []
        l_m = np.asarray(color_c)

        if (len(active) == 0) and (len(tbc) == 0):
            # this first case basically consist with the initial situation
            # if both active and tbc are emptied, all the new detections have to go into
            # tbc

            # the dection is done with a nested for loop, every element is compared to every other.
            # if there is a match within a certain threshold, then there is merge
            for i in range(0,np.shape(l_m)[0]):
                if (l_m[i][0] >= 0):
                    sum_x = l_m[i][0]
                    sum_y = l_m[i][1]
                    sum_r = l_m[i][2]
                    # r = l_m[i][2]
                    n = 1
                    for j in range(0, np.shape(l_m)[0]):
                        # I don't want to compare the element to itself, nor to an
                        # element which has already been assigned.
                        if(i != j and l_m[j][0]>= 0):
                            r = max(l_m[i][2],l_m[j][2])
                            dis1 = depth_image[int(l_m[i][1]),int(l_m[i][0])]*depth_scale
                            dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                            if(abs(dis1-dis2)<1.2*r*depth_scale):
                                if sqrt((l_m[i][0]-l_m[j][0])**2+(l_m[i][1]-l_m[j][1])**2) <= 1.2*r:
                                    # dis1 = depth_image[int(l_m[i][1]),int(l_m[i][0])]*depth_scale
                                    # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                                    # if (abs(dis1-dis2)<r):
                                    n = n + 1
                                    sum_x = sum_x + l_m[j][0]
                                    sum_y = sum_y + l_m[j][1]
                                    sum_r = sum_r + l_m[j][2]
                                    # print(type(l_m))
                                    l_m[j][0] = -1
                                    l_m[j][1] = -1
                    cX = int(round(sum_x/n))
                    cY = int(round(sum_y/n))
                    rf = int(round(sum_r/n))
                    tbc.append((cX,cY,rf))
        elif(len(active) == 0) and (len(tbc) != 0):
            # here now tbc has to be compared with the incoming measurements
            t_b = np.asarray(tbc)
            # empty the list, it will be repopulated anyways
            tbc *= 0
            for i in range(0,np.shape(t_b)[0]):
                flag = 0
                sum_x = t_b[i][0]
                sum_y = t_b[i][1]
                sum_r = t_b[i][2]
                # r = t_b[i][2]
                n = 1
                for j in range(0,np.shape(l_m)[0]):
                    # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
                    if (l_m[j][0]>0):
                        r = max(t_b[i][2],l_m[j][2])
                        dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
                        dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                        if(abs(dis1-dis2)<1.2*r*depth_scale):
                            if sqrt((t_b[i][0]-l_m[j][0])**2+(t_b[i][1]-l_m[j][1])**2) <= 1.2*r:
                                # dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
                                # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                                # if (abs(dis1-dis2)<r):
                                flag = 1
                                sum_x += l_m[j][0]
                                sum_y += l_m[j][1]
                                sum_r += l_m[j][2]
                                n += 1
                                l_m[j][0] = -1
                                l_m[j][1] = -1
                if(flag):
                    # only I have something matches I add it to the list of activte nodes
                    cX = int(round(sum_x/n))
                    cY = int(round(sum_y/n))
                    rf = int(round(sum_r/n))
                    active.append((cX,cY,rf))
            # I don't really like this part
            for k in range(0,np.shape(l_m)[0]):
                if(l_m[k][0] > 0):
                    tbc.append((l_m[k][0],l_m[k][1],l_m[k][2]))

        elif(len(active)!=0):
        # else:
            # this is the last case, in which we should have both tbc, active and the new incoming ones
            t_b = np.asarray(tbc)
            a = np.asarray(active)
            active *= 0
            tbc *= 0
            # I basically do what I did before. Fist the incoming measurments are compared with the active ones. If there are matches the active node is updated
            # otherwise the active nodes will be discarded
            for i in range(0,np.shape(a)[0]):
                flag = 0
                sum_x = a[i][0]
                sum_y = a[i][1]
                sum_r = a[i][2]
                r = a[i][2]
                n = 1
                for j in range(0,np.shape(l_m)[0]):
                    # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
                    if (l_m[j][0]>0):
                        r = max(a[i][2],l_m[j][2])
                        dis1 = depth_image[int(a[i][1]),int(a[i][0])]*depth_scale
                        dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                        if(abs(dis1-dis2)<1.2*r*depth_scale):
                            if sqrt((a[i][0]-l_m[j][0])**2+(a[i][1]-l_m[j][1])**2) <= 1.2*r:
                            #     dis1 = depth_image[int(a[i][1]),int(a[i][0])]*depth_scale
                            #     dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                                # if (abs(dis1-dis2)<r):
                                flag = 1
                                sum_x += l_m[j][0]
                                sum_y += l_m[j][1]
                                sum_r += l_m[j][2]
                                n += 1
                                l_m[j][0] = -1
                                l_m[j][1] = -1
                if(flag):
                    # only I have something matches I add it to the list of activte nodes
                    cX = int(round(sum_x/n))
                    cY = int(round(sum_y/n))
                    rf = int(round(sum_r/n))
                    active.append((cX,cY,rf))
            # After comparing to the active ones, I do the comparison to the tbc centers, If there are any matches
            # I add to active. If there is no match then I simply add the new detection to tbc
            for i in range(0,np.shape(t_b)[0]):
                flag = 0
                sum_x = t_b[i][0]
                sum_y = t_b[i][1]
                sum_r = t_b[i][2]
                r = t_b[i][2]
                n = 1
                for j in range(0,np.shape(l_m)[0]):
                    # yeah I cannot have multiple overlaps and like this I won't even be doing anything if a node was already considered before
                    if (l_m[j][0]>0):
                        r = max(t_b[i][2],l_m[j][2])
                        dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
                        dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                        if(abs(dis1-dis2)<1.2*r*depth_scale):
                            if sqrt((t_b[i][0]-l_m[j][0])**2+(t_b[i][1]-l_m[j][1])**2) <= 1.2*r:
                                # dis1 = depth_image[int(t_b[i][1]),int(t_b[i][0])]*depth_scale
                                # dis2 = depth_image[int(l_m[j][1]),int(l_m[j][0])]*depth_scale
                                # if (abs(dis1-dis2)<r):
                                flag = 1
                                sum_x += l_m[j][0]
                                sum_y += l_m[j][1]
                                sum_r += l_m[j][2]
                                n += 1
                                l_m[j][0] = -1
                                l_m[j][1] = -1
                if(flag):
                    # only I have something matches I add it to the list of activte nodes
                    cX = int(round(sum_x/n))
                    cY = int(round(sum_y/n))
                    rf = int(round(sum_r/n))
                    active.append((cX,cY,rf))
            for k in range(0,np.shape(l_m)[0]):
                if(l_m[k][0] > 0):
                    tbc.append((l_m[k][0],l_m[k][1],l_m[k][2]))
        red = []
        yellow = []
        pink = []
        black = []
        white = []
        blue = []
        green = []
        brown = []
        # this is what is needed in order to check whether we are dealing with a ball which is on the table
        lower = np.asarray([60,100,100])
        upper = np.asarray([80,255,255])
        mask = cv2.inRange(hsv, lower, upper)
        border = cv2.threshold(mask,150,255,cv2.THRESH_BINARY)[1]

        for z in range(0,np.shape(active)[0]):
                cX, cY, r = active[z]
                dis = depth_image[int(cY)-1,int(cX)]*depth_scale
                dis = round(dis,4)
                # check whether a minimum distance is fulfilled
                if(dis < 5.2 and dis> 0.2):
                    r_new = int(round(r*2))
                    cX_ma = cX + r_new
                    cX_mi = cX - r_new
                    cY_ma = cY + r_new
                    cY_mi = cY - r_new
                    if(cX_ma > 1279):
                        cX_ma = 1279
                    if(cX_mi < 0):
                        cX_mi = 0
                    if(cY_ma > 719):
                        cY_ma = 719
                    if(cY_mi < 0):
                        cY_mi = 0
                    # check if the ball is on the table
                    if(((border[cY_mi,cX]!=0)) or ((border[cY_ma,cX]!=0)) or ((border[cY,cX_mi]!=0)) or ((border[cY,cX_ma])!=0) ):
                      if np.sum(thr[0,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+10)))])/255 > 40 :
                          yellow.append((active[z]))
                      elif np.sum(thr[2,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          green.append((active[z]))
                      elif np.sum(thr[1,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          blue.append((active[z]))
                      elif np.sum(thr[3,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          red.append((active[z]))
                      elif( np.sum(thr[6,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                              int(round((active[z][0]-15))):int(round((active[z][0]+15)))])/255 >40):
                          pink.append((active[z]))
                      elif np.sum(thr[4,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                    (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          brown.append((active[z]))
                      elif np.sum(thr[5,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          white.append((active[z]))
                      elif np.sum(thr[7,int(round((active[z][1]-15))):int(round((active[z][1]+15))),\
                                  (int(round(active[z][0]-15))):int(round((active[z][0]+15)))])/255 > 40 :
                          black.append((active[z]))
        all = red + yellow + blue + green +brown + white + pink + black
        active *= 0
        active = list(all)
        # print(brown)
        t = np.asarray(all)
        # t = np.asarray(red)

        # print(t)
        # t = np.asarray(brown)
        # print('white')
        # print(white)
        # print('blue')
        # print(blue)


        # cv2.imshow('m', mask)

        # print(border[100,200])
        # print('start')
        for n in range(0, np.shape(t)[0]):
            cX, cY, r = t[n]
            dis = depth_image[int(cY),int(cX)]*depth_scale
            dis = round(dis,4)
            r_new = int(round(r*1.5))
            cX_ma = cX + r_new
            cX_mi = cX - r_new
            cY_ma = cY + r_new
            cY_mi = cY - r_new
            cv2.rectangle(color_image,(cX-r_new,cY-r_new),(cX+r_new,cY+r_new),(0, 255, 0), 1)
            cv2.putText(color_image, "Distance: " + str(dis), (int(cX)-r_new,int(cY)+2*r_new), font, 1, (0,0,255),2,cv2.LINE_AA)


        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET)
        cv2.imwrite("distance_image.png",color_image)
        # cv2.imwrite("dept_image.png", depth_colormap)
        images = np.hstack((color_image, depth_colormap))
        cv2.namedWindow('Align Example', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('Align Example', images)
        # cv2.imshow('red2', mask2)
        cv2.imshow('thr',thr[7,:,:])
        cv2.imshow('b',border)
        # cv2.imshow('Gray',gray)
        key = cv2.waitKey(1)
        # Press esc or 'q' to close the image window
        if key & 0xFF == ord('q') or key == 27:
            cv2.destroyAllWindows()
            break
finally:
    pipeline.stop()