from builtins import dict
from common.utils.global_variables import *

MIN_IOU_HEIGHT = 0.7
MIN_WIDTH_LINE_RATIO = 0.05


class Word:
    def __init__(
        self,
        text="",
        image=None,
        conf_detect=0.0,
        conf_cls=0.0,
        bndbox=None,
        kie_label="",
    ):
        self.type = "word"
        self.text = text
        self.image = image
        self.conf_detect = conf_detect
        self.conf_cls = conf_cls
        self.boundingbox = bndbox if bndbox is not None else [-1, -1, -1, -1]# [left, top,right,bot] coordinate of top-left and bottom-right point
        self.word_id = 0  # id of word
        self.word_group_id = 0  # id of word_group which instance belongs to
        self.line_id = 0  # id of line which instance belongs to
        self.paragraph_id = 0  # id of line which instance belongs to
        self.kie_label = kie_label

    def invalid_size(self):
        return (self.boundingbox[2] - self.boundingbox[0]) * (
            self.boundingbox[3] - self.boundingbox[1]
        ) > 0

    def is_special_word(self):
        left, top, right, bottom = self.boundingbox
        width, height = right - left, bottom - top
        text = self.text

        if text is None:
            return True

        if len(text) >= 7:
            no_digits = sum(c.isdigit() for c in text)
            return no_digits / len(text) >= 0.3

        return False


class Word_group:
    def __init__(self):
        self.type = "word_group"
        self.list_words = []  # dict of word instances
        self.word_group_id = 0  # word group id
        self.line_id = 0  # id of line which instance belongs to
        self.paragraph_id = 0  # id of paragraph which instance belongs to
        self.text = ""
        self.boundingbox = [-1, -1, -1, -1]
        self.kie_label = ""

    def add_word(self, word: Word):  # add a word instance to the word_group
        if word.text != "✪":
            for w in self.list_words:
                if word.word_id == w.word_id:
                    print("Word id collision")
                    return False
            word.word_group_id = self.word_group_id  #
            word.line_id = self.line_id
            word.paragraph_id = self.paragraph_id
            self.list_words.append(word)
            self.text += " " + word.text
            if self.boundingbox == [-1, -1, -1, -1]:
                self.boundingbox = word.boundingbox
            else:
                self.boundingbox = [
                    min(self.boundingbox[0], word.boundingbox[0]),
                    min(self.boundingbox[1], word.boundingbox[1]),
                    max(self.boundingbox[2], word.boundingbox[2]),
                    max(self.boundingbox[3], word.boundingbox[3]),
                ]
            return True
        else:
            return False

    def update_word_group_id(self, new_word_group_id):
        self.word_group_id = new_word_group_id
        for i in range(len(self.list_words)):
            self.list_words[i].word_group_id = new_word_group_id

    def update_kie_label(self):
        list_kie_label = [word.kie_label for word in self.list_words]
        dict_kie = dict()
        for label in list_kie_label:
            if label not in dict_kie:
                dict_kie[label] = 1
            else:
                dict_kie[label] += 1
        max_value = max(list(dict_kie.values()))
        list_keys = list(dict_kie.keys())
        list_values = list(dict_kie.values())
        self.kie_label = list_keys[list_values.index(max_value)]

    def update_text(self):  # update text after changing positions of words in list word
        text = ""
        for word in self.list_words:
            text += " " + word.text
        self.text = text


class Line:
    def __init__(self):
        self.type = "line"
        self.list_word_groups = []  # list of Word_group instances in the line
        self.line_id = 0  # id of line in the paragraph
        self.paragraph_id = 0  # id of paragraph which instance belongs to
        self.text = ""
        self.boundingbox = [-1, -1, -1, -1]

    def add_group(self, word_group: Word_group):  # add a word_group instance
        if word_group.list_words is not None:
            for wg in self.list_word_groups:
                if word_group.word_group_id == wg.word_group_id:
                    print("Word_group id collision")
                    return False

            self.list_word_groups.append(word_group)
            self.text += word_group.text
            word_group.paragraph_id = self.paragraph_id
            word_group.line_id = self.line_id

            for i in range(len(word_group.list_words)):
                word_group.list_words[
                    i
                ].paragraph_id = self.paragraph_id  # set paragraph_id for word
                word_group.list_words[i].line_id = self.line_id  # set line_id for word
            return True
        return False

    def update_line_id(self, new_line_id):
        self.line_id = new_line_id
        for i in range(len(self.list_word_groups)):
            self.list_word_groups[i].line_id = new_line_id
            for j in range(len(self.list_word_groups[i].list_words)):
                self.list_word_groups[i].list_words[j].line_id = new_line_id

    def merge_word(self, word):  # word can be a Word instance or a Word_group instance
        if word.text != "✪":
            if self.boundingbox == [-1, -1, -1, -1]:
                self.boundingbox = word.boundingbox
            else:
                self.boundingbox = [
                    min(self.boundingbox[0], word.boundingbox[0]),
                    min(self.boundingbox[1], word.boundingbox[1]),
                    max(self.boundingbox[2], word.boundingbox[2]),
                    max(self.boundingbox[3], word.boundingbox[3]),
                ]
            self.list_word_groups.append(word)
            self.text += " " + word.text
            return True
        return False

    def __cal_ratio(self, top1, bottom1, top2, bottom2):
        sorted_vals = sorted([top1, bottom1, top2, bottom2])
        intersection = sorted_vals[2] - sorted_vals[1]
        min_height = min(bottom1 - top1, bottom2 - top2)
        if min_height == 0:
            return -1
        ratio = intersection / min_height
        return ratio

    def __cal_ratio_height(self, top1, bottom1, top2, bottom2):

        height1, height2 = top1 - bottom1, top2 - bottom2
        ratio_height = float(max(height1, height2)) / float(min(height1, height2))
        return ratio_height

    def in_same_line(self, input_line, thresh=0.7):
        # calculate iou in vertical direction
        _, top1, _, bottom1 = self.boundingbox
        _, top2, _, bottom2 = input_line.boundingbox

        ratio = self.__cal_ratio(top1, bottom1, top2, bottom2)
        ratio_height = self.__cal_ratio_height(top1, bottom1, top2, bottom2)

        if (
            (top1 in range(top2, bottom2) or top2 in range(top1, bottom1))
            and ratio >= thresh
            and (ratio_height < 2)
        ):
            return True
        return False


class Paragraph:
    def __init__(self, id=0, lines=None):
        self.list_lines = (
            lines if lines is not None else []
        )  # list of all lines in the paragraph
        self.paragraph_id = id  # index of paragraph in the ist of paragraph
        self.text = ""
        self.boundingbox = [-1, -1, -1, -1]

    def add_line(self, line: Line):  # add a line instance
        if line.list_word_groups is not None:
            for l in self.list_lines:
                if line.line_id == l.line_id:
                    print("Line id collision")
                    return False
            for i in range(len(line.list_word_groups)):
                line.list_word_groups[
                    i
                ].paragraph_id = (
                    self.paragraph_id
                )  # set paragraph id for every word group in line
                for j in range(len(line.list_word_groups[i].list_words)):
                    line.list_word_groups[i].list_words[
                        j
                    ].paragraph_id = (
                        self.paragraph_id
                    )  # set paragraph id for every word in word groups
            line.paragraph_id = self.paragraph_id  # set paragraph id for line
            self.list_lines.append(line)  # add line to paragraph
            self.text += " " + line.text
            return True
        else:
            return False

    def update_paragraph_id(
        self, new_paragraph_id
    ):  # update new paragraph_id for all lines, word_groups, words inside paragraph
        self.paragraph_id = new_paragraph_id
        for i in range(len(self.list_lines)):
            self.list_lines[
                i
            ].paragraph_id = new_paragraph_id  # set new paragraph_id for line
            for j in range(len(self.list_lines[i].list_word_groups)):
                self.list_lines[i].list_word_groups[
                    j
                ].paragraph_id = new_paragraph_id  # set new paragraph_id for word_group
                for k in range(len(self.list_lines[i].list_word_groups[j].list_words)):
                    self.list_lines[i].list_word_groups[j].list_words[
                        k
                    ].paragraph_id = new_paragraph_id  # set new paragraph id for word
        return True


def resize_to_original(
    boundingbox, scale
):  # resize coordinates to match size of original image
    left, top, right, bottom = boundingbox
    left *= scale[1]
    right *= scale[1]
    top *= scale[0]
    bottom *= scale[0]
    return [left, top, right, bottom]


def check_iomin(word: Word, word_group: Word_group):
    min_height = min(
        word.boundingbox[3] - word.boundingbox[1],
        word_group.boundingbox[3] - word_group.boundingbox[1],
    )
    intersect = min(word.boundingbox[3], word_group.boundingbox[3]) - max(
        word.boundingbox[1], word_group.boundingbox[1]
    )
    if intersect / min_height > 0.7:
        return True
    return False


def prepare_line(words):
    lines = []
    visited = [False] * len(words)
    for id_word, word in enumerate(words):
        if word.invalid_size() == 0:
            continue
        new_line = True
        for i in range(len(lines)):
            if (
                lines[i].in_same_line(word) and not visited[id_word]
            ):  # check if word is in the same line with lines[i]
                lines[i].merge_word(word)
                new_line = False
                visited[id_word] = True

        if new_line == True:
            new_line = Line()
            new_line.merge_word(word)
            lines.append(new_line)

    # print(len(lines))
    # sort line from top to bottom according top coordinate
    lines.sort(key=lambda x: x.boundingbox[1])
    return lines


def __create_word_group(word, word_group_id):
    new_word_group = Word_group()
    new_word_group.word_group_id = word_group_id
    new_word_group.add_word(word)

    return new_word_group


def __sort_line(line):
    line.list_word_groups.sort(
        key=lambda x: x.boundingbox[0]
    )  # sort word in lines from left to right

    return line


def __merge_text_for_line(line):
    line.text = ""
    for word in line.list_word_groups:
        line.text += " " + word.text

    return line


def __update_list_word_groups(line, word_group_id, word_id, line_width):

    old_list_word_group = line.list_word_groups
    list_word_groups = []

    inital_word_group = __create_word_group(old_list_word_group[0], word_group_id)
    old_list_word_group[0].word_id = word_id
    list_word_groups.append(inital_word_group)
    word_group_id += 1
    word_id += 1

    for word in old_list_word_group[1:]:
        check_word_group = True
        word.word_id = word_id
        word_id += 1

        if (
            (not list_word_groups[-1].text.endswith(":"))
            and (
                (word.boundingbox[0] - list_word_groups[-1].boundingbox[2]) / line_width
                < MIN_WIDTH_LINE_RATIO
            )
            and check_iomin(word, list_word_groups[-1])
        ):
            list_word_groups[-1].add_word(word)
            check_word_group = False

        if check_word_group:
            new_word_group = __create_word_group(word, word_group_id)
            list_word_groups.append(new_word_group)
            word_group_id += 1
    line.list_word_groups = list_word_groups
    return line, word_group_id, word_id


def construct_word_groups_in_each_line(lines):
    line_id = 0
    word_group_id = 0
    word_id = 0
    for i in range(len(lines)):
        if len(lines[i].list_word_groups) == 0:
            continue

        # left, top ,right, bottom
        line_width = lines[i].boundingbox[2] - lines[i].boundingbox[0]  # right - left

        lines[i] = __sort_line(lines[i])

        # update text for lines after sorting
        lines[i] = __merge_text_for_line(lines[i])

        lines[i], word_group_id, word_id = __update_list_word_groups(
            lines[i], word_group_id, word_id, line_width
        )
        lines[i].update_line_id(line_id)
        line_id += 1
    return lines


def words_to_lines(words, check_special_lines=True):  # words is list of Word instance
    # sort word by top
    words.sort(key=lambda x: (x.boundingbox[1], x.boundingbox[0]))
    number_of_word = len(words)
    # print(number_of_word)
    # sort list words to list lines, which have not contained word_group yet
    lines = prepare_line(words)

    # construct word_groups in each line
    lines = construct_word_groups_in_each_line(lines)
    return lines, number_of_word


def near(word_group1: Word_group, word_group2: Word_group):
    min_height = min(
        word_group1.boundingbox[3] - word_group1.boundingbox[1],
        word_group2.boundingbox[3] - word_group2.boundingbox[1],
    )
    overlap = min(word_group1.boundingbox[3], word_group2.boundingbox[3]) - max(
        word_group1.boundingbox[1], word_group2.boundingbox[1]
    )

    if overlap > 0:
        return True
    if abs(overlap / min_height) < 1.5:
        print("near enough", abs(overlap / min_height), overlap, min_height)
        return True
    return False


def calculate_iou_and_near(wg1: Word_group, wg2: Word_group):
    min_height = min(
        wg1.boundingbox[3] - wg1.boundingbox[1], wg2.boundingbox[3] - wg2.boundingbox[1]
    )
    overlap = min(wg1.boundingbox[3], wg2.boundingbox[3]) - max(
        wg1.boundingbox[1], wg2.boundingbox[1]
    )
    iou = overlap / min_height
    distance = min(
        abs(wg1.boundingbox[0] - wg2.boundingbox[2]),
        abs(wg1.boundingbox[2] - wg2.boundingbox[0]),
    )
    if iou > 0.7 and distance < 0.5 * (wg1.boundingboxp[2] - wg1.boundingbox[0]):
        return True
    return False


def construct_word_groups_to_kie_label(list_word_groups: list):
    kie_dict = dict()
    for wg in list_word_groups:
        if wg.kie_label == "other":
            continue
        if wg.kie_label not in kie_dict:
            kie_dict[wg.kie_label] = [wg]
        else:
            kie_dict[wg.kie_label].append(wg)

    new_dict = dict()
    for key, value in kie_dict.items():
        if len(value) == 1:
            new_dict[key] = value
            continue

        value.sort(key=lambda x: x.boundingbox[1])
        new_dict[key] = value
    return new_dict


def invoice_construct_word_groups_to_kie_label(list_word_groups: list):
    kie_dict = dict()

    for wg in list_word_groups:
        if wg.kie_label == "other":
            continue
        if wg.kie_label not in kie_dict:
            kie_dict[wg.kie_label] = [wg]
        else:
            kie_dict[wg.kie_label].append(wg)

    return kie_dict


def postprocess_total_value(kie_dict):
    if "total_in_words_value" not in kie_dict:
        return kie_dict

    for k, value in kie_dict.items():
        if k == "total_in_words_value":
            continue
        l = []
        for v in value:
            if v.boundingbox[3] <= kie_dict["total_in_words_value"][0].boundingbox[3]:
                l.append(v)

        if len(l) != 0:
            kie_dict[k] = l

    return kie_dict


def postprocess_tax_code_value(kie_dict):
    if "buyer_tax_code_value" in kie_dict or "seller_tax_code_value" not in kie_dict:
        return kie_dict

    kie_dict["buyer_tax_code_value"] = []
    for v in kie_dict["seller_tax_code_value"]:
        if "buyer_name_key" in kie_dict and (
            v.boundingbox[3] > kie_dict["buyer_name_key"][0].boundingbox[3]
            or near(v, kie_dict["buyer_name_key"][0])
        ):
            kie_dict["buyer_tax_code_value"].append(v)
            continue

        if "buyer_name_value" in kie_dict and (
            v.boundingbox[3] > kie_dict["buyer_name_value"][0].boundingbox[3]
            or near(v, kie_dict["buyer_name_value"][0])
        ):
            kie_dict["buyer_tax_code_value"].append(v)
            continue

        if "buyer_address_value" in kie_dict and near(
            kie_dict["buyer_address_value"][0], v
        ):
            kie_dict["buyer_tax_code_value"].append(v)
    return kie_dict


def postprocess_tax_code_key(kie_dict):
    if "buyer_tax_code_key" in kie_dict or "seller_tax_code_key" not in kie_dict:
        return kie_dict
    kie_dict["buyer_tax_code_key"] = []
    for v in kie_dict["seller_tax_code_key"]:
        if "buyer_name_key" in kie_dict and (
            v.boundingbox[3] > kie_dict["buyer_name_key"][0].boundingbox[3]
            or near(v, kie_dict["buyer_name_key"][0])
        ):
            kie_dict["buyer_tax_code_key"].append(v)
            continue

        if "buyer_name_value" in kie_dict and (
            v.boundingbox[3] > kie_dict["buyer_name_value"][0].boundingbox[3]
            or near(v, kie_dict["buyer_name_value"][0])
        ):
            kie_dict["buyer_tax_code_key"].append(v)
            continue

        if "buyer_address_value" in kie_dict and near(
            kie_dict["buyer_address_value"][0], v
        ):
            kie_dict["buyer_tax_code_key"].append(v)

    return kie_dict


def invoice_postprocess(kie_dict: dict):
    # all keys or values which are below total_in_words_value will be thrown away
    kie_dict = postprocess_total_value(kie_dict)
    kie_dict = postprocess_tax_code_value(kie_dict)
    kie_dict = postprocess_tax_code_key(kie_dict)
    return kie_dict


def throw_overlapping_words(list_words):
    new_list = [list_words[0]]
    for word in list_words:
        overlap = False
        area = (word.boundingbox[2] - word.boundingbox[0]) * (
            word.boundingbox[3] - word.boundingbox[1]
        )
        for word2 in new_list:
            area2 = (word2.boundingbox[2] - word2.boundingbox[0]) * (
                word2.boundingbox[3] - word2.boundingbox[1]
            )
            xmin_intersect = max(word.boundingbox[0], word2.boundingbox[0])
            xmax_intersect = min(word.boundingbox[2], word2.boundingbox[2])
            ymin_intersect = max(word.boundingbox[1], word2.boundingbox[1])
            ymax_intersect = min(word.boundingbox[3], word2.boundingbox[3])
            if xmax_intersect < xmin_intersect or ymax_intersect < ymin_intersect:
                continue

            area_intersect = (xmax_intersect - xmin_intersect) * (
                ymax_intersect - ymin_intersect
            )
            if area_intersect / area > 0.7 or area_intersect / area2 > 0.7:
                overlap = True
        if overlap == False:
            new_list.append(word)
    return new_list


class Box:
    def __init__(self, xmin=0, ymin=0, xmax=0, ymax=0, label="", kie_label=""):
        self.xmax = xmax
        self.ymax = ymax
        self.xmin = xmin
        self.ymin = ymin
        self.label = label
        self.kie_label = kie_label


def check_iou(box1: Word, box2: Box, threshold=0.9):
    area1 = (box1.boundingbox[2] - box1.boundingbox[0]) * (
        box1.boundingbox[3] - box1.boundingbox[1]
    )
    area2 = (box2.xmax - box2.xmin) * (box2.ymax - box2.ymin)
    xmin_intersect = max(box1.boundingbox[0], box2.xmin)
    ymin_intersect = max(box1.boundingbox[1], box2.ymin)
    xmax_intersect = min(box1.boundingbox[2], box2.xmax)
    ymax_intersect = min(box1.boundingbox[3], box2.ymax)
    if xmax_intersect < xmin_intersect or ymax_intersect < ymin_intersect:
        area_intersect = 0
    else:
        area_intersect = (xmax_intersect - xmin_intersect) * (
            ymax_intersect - ymin_intersect
        )
    union = area1 + area2 - area_intersect
    iou = area_intersect / union
    if iou > threshold:
        return True
    return False
