import cv2
import pytesseract
import numpy as np
from pprint import pprint
from PIL import Image

img = cv2.imread('/Users/cgranados/Code/terravision/img/floor.png')
grey = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
hsv = img

thresh_map = [
    {
        'name': 'Kitchen',
        'colors': {'low': [95, 199, 212], 'high': [115, 219, 292]},
        'max_area': 0
    },
    {
        'name': 'AV Room',
        'colors': {'low': [], 'high': []},
        'max_area': 0
    },
    {
        'name': 'Restroom',
        'colors': {'low': [], 'high': []},
        'max_area': 0
    },
    {
        'name': 'Private Office',
        'colors': {'low': [], 'high': []},
        'max_area': 0
    }
]


# HSV value 42, 58, 100
# Declare lower bound and upper bound for color
lower_bound = np.array(thresh_map[0]['colors']['low'])
upper_bound = np.array(thresh_map[0]['colors']['high'])

# lower_bound = np.array([245, 245, 197])
# upper_bound = np.array([265, 265, 277])

# Create mask with pixels found in range of upper and lower bound
mask = cv2.inRange(hsv, lower_bound, upper_bound)

# If area of interest if larger than we expect de-noise with erosion.
kernel = np.ones((11, 11), np.uint8)
erosion = cv2.erode(mask, kernel, iterations=1)

# Dilate to so we can recover the area on the larger image
kernel = np.ones((8, 8), np.uint8)
dilation = cv2.dilate(erosion, kernel, iterations=2)

# kernel = np.ones((1, 9), np.uint8)
# dilation = cv2.dilate(dilation, kernel, iterations=1)

# res = cv2.bitwise_and(img, img, mask=dilation)
#
# cv2.imshow('res', res)
#
# cv2.waitKey(0)

ret, thresh = cv2.threshold(grey, 127, 255, 0)
_, contours, hierarchy = cv2.findContours(dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

cnt = contours[0]
x, y, w, h = cv2.boundingRect(cnt)

crop = img[y:y+h, x:x+w]

cv2.imshow('original', img)

cv2.waitKey(0)

cv2.imshow('dialated', dilation)

cv2.waitKey(0)

cv2.imshow("final", crop)

cv2.waitKey(0)

test_mask = cv2.inRange(crop, np.array([50, 50, 50]), np.array([130, 130, 130]))

cv2.imshow('masked text', test_mask)

cv2.waitKey(0)

cv2.destroyAllWindows()

r = "/Users/cgranados/Code/terravision/img/floor_extracted.png"
cv2.imwrite(r, crop)

text = pytesseract.image_to_string(Image.open(r))

print "text found:" + text

cv2.destroyAllWindows()

exit()

cnt = contours[0]

M = cv2.moments(cnt)

height, width, channels = img.shape

new_mask = np.zeros((height, width), np.uint8)

points = np.array(contours)

# cv2.fillPoly(new_mask, pts=np.int32([points]), color=(255, 255, 255))

for c in contours:
    # Find the minimum bounding box without accounting for orientation
    x, y, w, h = cv2.boundingRect(c)
    # Draw rectangle on binary mask
    cv2.rectangle(new_mask, (x, y), (x+w, y+h), (255, 255, 255), -1)

# box = [[(x, y), (x + w, y + h)] for x, y, w, h in cv2.boundingRect(c)]

# Isolate area of interest with mask
res = cv2.bitwise_and(img, img, mask=new_mask)

cv2.imshow('res', res)

crop = img[y:y+h,x:x+w]

cv2.waitKey(0)

# cv2.drawContours(new_mask, box, )

# epsilon = 10 * cv2.arcLength(cnt, True)
# approx = cv2.approxPolyDP(cnt, epsilon, True)
#
# hull = cv2.convexHull(cnt)
#
# pprint(hull)
#
# pprint(cnt)
#
# exit(0)

# cv2.drawContours(img, [hull], 0, (0, 255, 0), 3)
#
# cv2.imshow('draw', img)
#
# cv2.waitKey(0)
#
# cv2.destroyAllWindows()

cv2.imshow('original', img)

cv2.waitKey(0)
#
# cv2.imshow('found', mask)
#
# cv2.waitKey(0)
#
# cv2.imshow('eroded', erosion)
#
# cv2.waitKey(0)
#
cv2.imshow('dialated', dilation)

cv2.waitKey(0)

cv2.imshow("new mask", new_mask)

cv2.waitKey(0)

cv2.destroyAllWindows()

# remove the background

# def remove_background(matrix, color, treshhold, exact=False):
#     '''
#     Slow as a dog but acceptable for now. Runs at a X*Y complexity.
#     :param matrix:
#     :return: matrix
#     '''
#     for y in matrix:
#         for x in y:
#             if x >=

# Algorithm

# Declare lower bound and upper bound for color
# Create mask with pixels found in range of upper and lower bound, returns bitwise
# If area of interest if larger than we expect de-noise with erosion.
# Dilate to so we can recover the area on the larger image
# Find contours in the shape
# Get convex hull to simplify into square
# Use result (contours) to draw a new square binary image
# Use binary image as mask on the original image, we now have the isolated room
# Get it's coordinates on a 2D plane.
# Find the shape's midpoint coordinate. (this is part of the info we want to return)
# Crop out all black from the image so we only have the room in the matrix (and we deal with a smaller matrix)
# Perform optical character recognition on isolated room image.
# Refine OCR because it won't be perfect:
#    - Iterate through space delimited tokens and try for exact matches against a dictionary of expected words.
#    - If no exact match then use trie data structure with indexed rooms for characters until we find a match.
#      Each time starting from the next character if no match found.
#
# Output:
#    - Coordinates of hull (square that surrounds the room) in 2D space, coordinates
#    - Coordinate of midpoint in 2D space.


# OCR refining algorithm.

# grab token and see if it's exact match.

# otherwise use trie structure and grab stream from start to end.