import math
import os
from collections import OrderedDict
from pathlib import Path
from typing import Optional, Union
from matplotlib import colors
import copy

import cv2
import numpy as np

import craft_text_detector.file_utils as file_utils
import craft_text_detector.torch_utils as torch_utils

# unwarp corodinates
def warpCoord(Minv, pt):
    out = np.matmul(Minv, (pt[0], pt[1], 1))
    return np.array([out[0] / out[2], out[1] / out[2]])

def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text):
    # prepare data
    linkmap = linkmap.copy()
    textmap = textmap.copy()
    img_h, img_w = textmap.shape

    """ labeling method """
    ret, text_score = cv2.threshold(textmap, low_text, 1, 0)
    ret, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)

    text_score_comb = np.clip(text_score + link_score, 0, 1)

    nLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        text_score_comb.astype(np.uint8), connectivity=4
    )

    det = []
    det_scores = []
    mapper = []
    for k in range(1, nLabels):
        # size filtering
        size = stats[k, cv2.CC_STAT_AREA]
        if size < 10:
            continue

        # thresholding
        if np.max(textmap[labels == k]) < text_threshold:
            continue

        # make segmentation map
        segmap = np.zeros(textmap.shape, dtype=np.uint8)
        segmap[labels == k] = 255

        # remove link area
        segmap[np.logical_and(link_score == 1, text_score == 0)] = 0

        x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
        w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
        niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
        sx, ex, sy, ey = (x - niter, x + w + niter + 1, y - niter, y + h + niter + 1)
        # boundary check
        if sx < 0:
            sx = 0
        if sy < 0:
            sy = 0
        if ex >= img_w:
            ex = img_w
        if ey >= img_h:
            ey = img_h
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
        segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)

        # make box
        np_temp = np.roll(np.array(np.where(segmap != 0)), 1, axis=0)
        np_contours = np_temp.transpose().reshape(-1, 2)
        rectangle = cv2.minAreaRect(np_contours)
        box = cv2.boxPoints(rectangle)

        # boundary check due to minAreaRect may have out of range values 
        # (see https://docs.opencv.org/3.4/d3/dc0/group__imgproc__shape.html#ga3d476a3417130ae5154aea421ca7ead9)
        for p in box:
            if p[0] < 0:
                p[0] = 0
            if p[1] < 0:
                p[1] = 0
            if p[0] >= img_w:
                p[0] = img_w
            if p[1] >= img_h:
                p[1] = img_h

        # align diamond-shape
        w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
        box_ratio = max(w, h) / (min(w, h) + 1e-5)
        if abs(1 - box_ratio) <= 0.1:
            l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
            t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
            box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)

        # make clock-wise order
        startidx = box.sum(axis=1).argmin()
        box = np.roll(box, 4 - startidx, 0)
        box = np.array(box)

        det.append(box)
        mapper.append(k)
        det_scores.append(np.max(textmap[labels==k]))

    return det, labels, mapper, det_scores


def getPoly_core(boxes, labels, mapper, linkmap):
    # configs
    num_cp = 5
    max_len_ratio = 0.7
    expand_ratio = 1.45
    max_r = 2.0
    step_r = 0.2

    polys = []
    for k, box in enumerate(boxes):
        # size filter for small instance
        w, h = (
            int(np.linalg.norm(box[0] - box[1]) + 1),
            int(np.linalg.norm(box[1] - box[2]) + 1),
        )
        if w < 10 or h < 10:
            polys.append(None)
            continue

        # warp image
        tar = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
        M = cv2.getPerspectiveTransform(box, tar)
        word_label = cv2.warpPerspective(labels, M, (w, h), flags=cv2.INTER_NEAREST)
        try:
            Minv = np.linalg.inv(M)
        except:
            polys.append(None)
            continue

        # binarization for selected label
        cur_label = mapper[k]
        word_label[word_label != cur_label] = 0
        word_label[word_label > 0] = 1

        """ Polygon generation """
        # find top/bottom contours
        cp = []
        max_len = -1
        for i in range(w):
            region = np.where(word_label[:, i] != 0)[0]
            if len(region) < 2:
                continue
            cp.append((i, region[0], region[-1]))
            length = region[-1] - region[0] + 1
            if length > max_len:
                max_len = length

        # pass if max_len is similar to h
        if h * max_len_ratio < max_len:
            polys.append(None)
            continue

        # get pivot points with fixed length
        tot_seg = num_cp * 2 + 1
        seg_w = w / tot_seg  # segment width
        pp = [None] * num_cp  # init pivot points
        cp_section = [[0, 0]] * tot_seg
        seg_height = [0] * num_cp
        seg_num = 0
        num_sec = 0
        prev_h = -1
        for i in range(0, len(cp)):
            (x, sy, ey) = cp[i]
            if (seg_num + 1) * seg_w <= x and seg_num <= tot_seg:
                # average previous segment
                if num_sec == 0:
                    break
                cp_section[seg_num] = [
                    cp_section[seg_num][0] / num_sec,
                    cp_section[seg_num][1] / num_sec,
                ]
                num_sec = 0

                # reset variables
                seg_num += 1
                prev_h = -1

            # accumulate center points
            cy = (sy + ey) * 0.5
            cur_h = ey - sy + 1
            cp_section[seg_num] = [
                cp_section[seg_num][0] + x,
                cp_section[seg_num][1] + cy,
            ]
            num_sec += 1

            if seg_num % 2 == 0:
                continue  # No polygon area

            if prev_h < cur_h:
                pp[int((seg_num - 1) / 2)] = (x, cy)
                seg_height[int((seg_num - 1) / 2)] = cur_h
                prev_h = cur_h

        # processing last segment
        if num_sec != 0:
            cp_section[-1] = [cp_section[-1][0] / num_sec, cp_section[-1][1] / num_sec]

        # pass if num of pivots is not sufficient or segment widh
        # is smaller than character height
        if None in pp or seg_w < np.max(seg_height) * 0.25:
            polys.append(None)
            continue

        # calc median maximum of pivot points
        half_char_h = np.median(seg_height) * expand_ratio / 2

        # calc gradiant and apply to make horizontal pivots
        new_pp = []
        for i, (x, cy) in enumerate(pp):
            dx = cp_section[i * 2 + 2][0] - cp_section[i * 2][0]
            dy = cp_section[i * 2 + 2][1] - cp_section[i * 2][1]
            if dx == 0:  # gradient if zero
                new_pp.append([x, cy - half_char_h, x, cy + half_char_h])
                continue
            rad = -math.atan2(dy, dx)
            c, s = half_char_h * math.cos(rad), half_char_h * math.sin(rad)
            new_pp.append([x - s, cy - c, x + s, cy + c])

        # get edge points to cover character heatmaps
        isSppFound, isEppFound = False, False
        grad_s = (pp[1][1] - pp[0][1]) / (pp[1][0] - pp[0][0]) + (
            pp[2][1] - pp[1][1]
        ) / (pp[2][0] - pp[1][0])
        grad_e = (pp[-2][1] - pp[-1][1]) / (pp[-2][0] - pp[-1][0]) + (
            pp[-3][1] - pp[-2][1]
        ) / (pp[-3][0] - pp[-2][0])
        for r in np.arange(0.5, max_r, step_r):
            dx = 2 * half_char_h * r
            if not isSppFound:
                line_img = np.zeros(word_label.shape, dtype=np.uint8)
                dy = grad_s * dx
                p = np.array(new_pp[0]) - np.array([dx, dy, dx, dy])
                cv2.line(
                    line_img,
                    (int(p[0]), int(p[1])),
                    (int(p[2]), int(p[3])),
                    1,
                    thickness=1,
                )
                if (
                    np.sum(np.logical_and(word_label, line_img)) == 0
                    or r + 2 * step_r >= max_r
                ):
                    spp = p
                    isSppFound = True
            if not isEppFound:
                line_img = np.zeros(word_label.shape, dtype=np.uint8)
                dy = grad_e * dx
                p = np.array(new_pp[-1]) + np.array([dx, dy, dx, dy])
                cv2.line(
                    line_img,
                    (int(p[0]), int(p[1])),
                    (int(p[2]), int(p[3])),
                    1,
                    thickness=1,
                )
                if (
                    np.sum(np.logical_and(word_label, line_img)) == 0
                    or r + 2 * step_r >= max_r
                ):
                    epp = p
                    isEppFound = True
            if isSppFound and isEppFound:
                break

        # pass if boundary of polygon is not found
        if not (isSppFound and isEppFound):
            polys.append(None)
            continue

        # make final polygon
        poly = []
        poly.append(warpCoord(Minv, (spp[0], spp[1])))
        for p in new_pp:
            poly.append(warpCoord(Minv, (p[0], p[1])))
        poly.append(warpCoord(Minv, (epp[0], epp[1])))
        poly.append(warpCoord(Minv, (epp[2], epp[3])))
        for p in reversed(new_pp):
            poly.append(warpCoord(Minv, (p[2], p[3])))
        poly.append(warpCoord(Minv, (spp[2], spp[3])))

        # add to final result
        polys.append(np.array(poly))

    return polys


def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False):
    boxes, labels, mapper, det_scores = getDetBoxes_core(
        textmap, linkmap, text_threshold, link_threshold, low_text
    )

    if poly:
        polys = getPoly_core(boxes, labels, mapper, linkmap)
    else:
        polys = [None] * len(boxes)

    return boxes, polys, det_scores


def adjustResultCoordinates(polys, ratio_w, ratio_h, ratio_net=2):
    if len(polys) > 0:
        polys = np.array(polys)
        for k in range(len(polys)):
            if polys[k] is not None:
                polys[k] *= (ratio_w * ratio_net, ratio_h * ratio_net)
    return polys


def post_process_regions(image, regions, inds, light_color, dark_color, inverse):
    # isolate blue color in HSV Color Space, and then isolate the green digital screen and crop only that

    region_white_pix = []
    for i in inds:
        # rectify and crop poly region
        crop_rgb = file_utils.crop_poly(image, regions[i])

        # normalize
        pixel_colors = crop_rgb.reshape((np.shape(crop_rgb)[0]*np.shape(crop_rgb)[1], 3))
        norm = colors.Normalize(vmin=-1.,vmax=1.)
        norm.autoscale(pixel_colors)
        pixel_colors = norm(pixel_colors).tolist()

        # converting image to HSV Color Space
        hsv_img = cv2.cvtColor(crop_rgb, cv2.COLOR_RGB2HSV)

        # picking out blue range (H,S,V)
        mask = cv2.inRange(hsv_img, light_color, dark_color)

        # Perform morphological operations
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
        opening = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        close = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel, iterations=1)

        if inverse:
            close[close == 0] = 150
            close[close == 255] = 0
            close[close == 150] = 255

        # Find contours and filter for largest contour
        # Draw largest contour onto a blank mask then bitwise-and
        blank_mask = np.zeros(crop_rgb.shape, dtype=np.uint8)
        cnts = cv2.findContours(close, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cnts = cnts[0] if len(cnts) == 2 else cnts[1]
        if cnts:
            cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[0]
            cv2.fillPoly(blank_mask, [cnts], (255,255,255))
            blank_mask = cv2.cvtColor(blank_mask, cv2.COLOR_BGR2GRAY)

            # count number of white (blue) pixels
            region_white_pix.append(np.sum(blank_mask == 255))

            # Crop ROI from result
            x,y,w,h = cv2.boundingRect(blank_mask)
            
            # update respective region
            start_left_width = min(regions[i][0][0],regions[i][3][0])
            regions[i][0][0] = start_left_width + x
            regions[i][3][0] = start_left_width + x
            regions[i][1][0] = regions[i][0][0] + w
            regions[i][2][0] = regions[i][3][0] + w

            start_top_height = min(regions[i][0][1],regions[i][1][1])
            regions[i][0][1] = start_top_height + y
            regions[i][2][1] = regions[i][0][1] + h
            regions[i][1][1] = regions[i][0][1]
            regions[i][3][1] = regions[i][0][1] + h

        else:
            region_white_pix.append(0)

    return regions, region_white_pix


def find_big_areas_poly(regions, n=1):
    """
    Arguments:
        regions: list of bboxes or polys
    """

    idx = None
    areas = []
    for region in regions:
        # region list is defined as: [[TL], [TR], [BR], [BL]]
        areas.append(abs(region[0][0]-region[1][0]) * abs(region[0][1]-region[3][1]))

    n = min(len(areas), n)
    idx = [areas.index(x) for x in sorted(areas, reverse=True)[:n]]

    return idx

def export_detected_regions(
    image,
    regions,
    regions_as_ratios,
    TL_TR_pixel_sub,
    file_name: str = "image",
):
    """
    Arguments:
        image: path to the image to be processed or numpy array or PIL image
        regions: list of bboxes or polys
        file_name (str): export image file name
        output_dir: folder to be exported
        rectify: rectify detected polygon by affine transform
    """

    # deepcopy image so that original is not altered
    image = copy.deepcopy(image)

    # exclude regions with height ratios bigger than width ratios
    new_regions = []
    valid_inds = []
    for ind, region in enumerate(regions):
        x_distance = ((region[0][0] - region[1][0])**2 + (region[0][1] - region[1][1])**2)**0.5
        y_distance = ((region[1][0] - region[2][0])**2 + (region[1][1] - region[2][1])**2)**0.5
        if x_distance > y_distance:
            new_regions.append(region)
            valid_inds.append(ind)

    regions_as_ratios = [e for i,e in enumerate(regions_as_ratios) if i in valid_inds]
    regions = [e for i,e in enumerate(regions) if i in valid_inds]

    # keep track of the original regions[ind] without HSV postprocessing
    # regions_ind_orig = regions
    orig_left_border = []
    for i, region in enumerate(regions):
        orig_left_border.append([regions[i][0][0], regions[i][3][0]])

    # Post-process detected regions
    # extend boundaries of regions first two inds in all directions to make sure we include the blue box with screen in the image
    # check what's the region with biggest amount of blue and keep that one
    # region list is defined as: [[TL], [TR], [BR], [BL]]
    right_offset = 100      # pixels
    left_offset = 100       # pixels
    height_offset = 75     # 75 pixels

    for i,_ in enumerate(regions):
        regions[i][1][0] = min(image.shape[1], regions[i][1][0] + right_offset)
        regions[i][2][0] = min(image.shape[1], regions[i][2][0] + right_offset)
        regions[i][0][0] = max(0, regions[i][0][0] - left_offset)
        regions[i][3][0] = max(0, regions[i][3][0] - left_offset)

        regions[i][1][1] = max(0, regions[i][1][1] - height_offset)
        regions[i][2][1] = min(image.shape[0], regions[i][2][1] + height_offset)
        regions[i][0][1] = max(0, regions[i][0][1] - height_offset)
        regions[i][3][1] = min(image.shape[0], regions[i][3][1] + left_offset)

    # find the 3 regions indexes with biggest area (sorted from biggest to lowest) from current regions
    inds = find_big_areas_poly(regions, len(regions))

    # isolate blue color in HSV Color Space, and then isolate the green digital screen and crop only that
    # region_white_pixels as the same item order as inds (sorted from biggest to lowest) !
    light_color = (75, 50, 100)
    dark_color = (140, 255, 255)
    regions, region_white_pixels = post_process_regions(image, regions, inds, light_color, dark_color, False)

    if not inds:
        ind = None
    else:
        # if the number of blue pixels for region 2 is superior to region, then use region 1 ind instead
        max_index = region_white_pixels.index(max(region_white_pixels))
        if len(inds) >= 2 and region_white_pixels[max_index] > region_white_pixels[0]:
            ind = max_index
        else:
            ind = inds[0]

        # rerun preprocessing to detect blue part of monitor, but this time delete it in order to keep the green screen
        regions, _ = post_process_regions(image, regions, inds, light_color, dark_color, True)

        # eliminate 32 % height pixels of bottom (BR and BL) digital traces assuming a fixed ratio of 0.68
        regions[ind][2][1] = regions[ind][3][1] = regions[ind][2][1] - (regions[ind][3][1]-regions[ind][0][1])*0.32

        # given the whole green digital monitor, remove a fixed 0.365 (=1-0.635) width ration of pixels (without chars) on the left width (TL and BL)
        regions[ind][0][0] = regions[ind][3][0] = regions[ind][0][0] + (regions[ind][1][0] - regions[ind][0][0])*0.365

        # subtracting TL_TR_pixel_sub pixels (works well) from the TL and TR height [1] to make the top of the numbers fully visiable and
        # maximize performance of PARSeq models
        regions[ind][0][1] = regions[ind][1][1] = regions[ind][0][1] - TL_TR_pixel_sub

    if ind is None:
        return None
    else:
        return regions[ind]
        