import pytesseract as pyt
import cv2

import lib_ip.ip_draw as draw
from config.CONFIG_UIED import Config

C = Config()


def is_text(img, min_word_area, show=False):
    broad = img.copy()
    area_word = 0
    area_total = img.shape[0] * img.shape[1]

    try:
        # ocr text detection
        data = pyt.image_to_data(img).split('\n')
    except:
        print(img.shape)
        return -1
    word = []
    for d in data[1:]:
        d = d.split()
        if d[-1] != '-1':
            if d[-1] != '-' and d[-1] != '—' and int(d[-3]) < 50 and int(d[-4]) < 100:
                word.append(d)
                t_l = (int(d[-6]), int(d[-5]))
                b_r = (int(d[-6]) + int(d[-4]), int(d[-5]) + int(d[-3]))
                area_word += int(d[-4]) * int(d[-3])
                cv2.rectangle(broad, t_l, b_r, (0,0,255), 1)

    if show:
        for d in word: print(d)
        print(area_word/area_total)
        cv2.imshow('a', broad)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    # no text in this clip or relatively small text area
    if len(word) == 0 or area_word/area_total < min_word_area:
        return False
    return True


def text_detection(org, img_clean):
    try:
        data = pyt.image_to_data(img_clean).split('\n')
    except:
        return org, None
    corners_word = []
    for d in data[1:]:
        d = d.split()
        if d[-1] != '-1':
            if d[-1] != '-' and d[-1] != '—' and 5 < int(d[-3]) < 40 and 5 < int(d[-4]) < 100:
                t_l = (int(d[-6]), int(d[-5]))
                b_r = (int(d[-6]) + int(d[-4]), int(d[-5]) + int(d[-3]))
                corners_word.append((t_l, b_r))
    return corners_word


# def text_merge_word_into_line(org, corners_word, max_words_gap=C.THRESHOLD_TEXT_MAX_WORD_GAP):
#
#     def is_in_line(word):
#         for i in range(len(lines)):
#             line = lines[i]
#             # at the same row
#             if abs(line['center'][1] - word['center'][1]) < max_words_gap:
#                 # small gap between words
#                 if (abs(line['center'][0] - word['center'][0]) - abs(line['width']/2 + word['width']/2)) < max_words_gap:
#                     return i
#         return -1
#
#     def merge_line(word, index):
#         line = lines[index]
#         # on the left
#         if word['center'][0] < line['center'][0]:
#             line['col_min'] = word['col_min']
#         # on the right
#         else:
#             line['col_max'] = word['col_max']
#         line['row_min'] = min(line['row_min'], word['row_min'])
#         line['row_max'] = max(line['row_max'], word['row_max'])
#         line['width'] = line['col_max'] - line['col_min']
#         line['height'] = line['row_max'] - line['row_min']
#         line['center'] = ((line['col_max'] + line['col_min'])/2, (line['row_max'] + line['row_min'])/2)
#
#     words = []
#     for corner in corners_word:
#         word = {}
#         (top_left, bottom_right) = corner
#         (col_min, row_min) = top_left
#         (col_max, row_max) = bottom_right
#         word['col_min'], word['col_max'], word['row_min'], word['row_max'] = col_min, col_max, row_min, row_max
#         word['height'] = row_max - row_min
#         word['width'] = col_max - col_min
#         word['center'] = ((col_max + col_min)/2, (row_max + row_min)/2)
#         words.append(word)
#
#     lines = []
#     for word in words:
#         line_index = is_in_line(word)
#         # word is in current line
#         if line_index != -1:
#             merge_line(word, line_index)
#         # word is not in current line
#         else:
#             # this single word as a new line
#             lines.append(word)
#
#     corners_line = []
#     for l in lines:
#         corners_line.append(((l['col_min'], l['row_min']), (l['col_max'], l['row_max'])))
#     return corners_line