import re

import numpy as np

from sdsvkie.utils.post_processing import near
from sdsvkie.utils.post_processing.common_post_processing import normalize_number_wordgroup, normalize_tax_wordgroup
from sdsvkie.utils.post_processing.post_processing_datetime import post_processing_datetime
from sdsvkie.utils import merge_boxes, Word_group
from sdsvkie.utils.post_processing.receipt_post_processing import is_wg_same_line

# 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 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 format_output(kie_dict):
        """
        {
                'field_key': {
                    'value': str,
                    'confidence': float,
                    'box': list[int]
                },
                or 'field_key':  {
                    'value': "",
                    'confidence': None,
                    'box': []
                },
            }

        Args:
            kie_dict (_type_): _description_

        Returns:
            _type_: _description_
        """
        end2end_results = {}
        for field_name, wg_list in kie_dict.items():
            if "key" in field_name:
                continue

            if not isinstance(wg_list, list):
                wg_list = [wg_list]

            if len(wg_list) == 0:
                text = ""
                conf = 0.0
                box = []
            else:
                
                text = " ".join([wg.text for wg in wg_list])
                text = text.strip()
                text = text.replace("✪", " ")
                conf_list = [wg.conf for wg in wg_list]
                conf = sum(conf_list) / len(conf_list)
                box_list = [wg.boundingbox for wg in wg_list]
                box = merge_boxes(box_list)

            end2end_results[field_name] = {
                "box": box,
                "value": text, 
                "conf": conf
            }

        # sort by keys
        field_keys = list(end2end_results.keys())
        field_keys.sort()
        end2end_results = {
            field_key: end2end_results[field_key] for field_key in field_keys
        }
        return end2end_results


def postprocess_date(date_wordgroups):
    date_wordgroups = sorted(date_wordgroups, key=lambda date_wg: date_wg.boundingbox[1]) 
    
    # merge total value if have multi boxes in one line
    line = [date_wordgroups[0]]
    for idx, wg in enumerate(date_wordgroups[1:]):
        if is_wg_same_line(line[-1], wg):
            line.append(wg)
        
    date_line = sorted(line, key=lambda wg: wg.boundingbox[0])   # sort wg by x-axis in a line
    if len(date_line) == 1:
        date_wg = date_line[0]
        date_wg.text = post_processing_datetime(date_wg.text)
        return [date_wg]
    
    #text
    date_text = " ".join([wg.text for wg in date_line if wg.text.strip() != ""])
    date_text = post_processing_datetime(date_text)
    #box
    date_box = merge_boxes([wg.boundingbox for wg in date_line])
    #label
    date_label = date_line[0].kie_label
    #conf
    conf_list = [wg.conf for wg in date_line]
    date_conf = sum(conf_list) / len(conf_list)

    # new value wg
    date_wg = Word_group()
    date_wg._set_values(
        boundingbox=date_box,
        text=date_text,
        kie_label=date_label,
        conf=date_conf
    )


    return [date_wg]

# def postprocess_total_in_words(total_words_wgs):
#     total_words_wgs = [normalize_total_words(total_words_wg) for total_words_wg in total_words_wgs]
#     return total_words_wgs



def invoice_postprocess(kie_dict: dict):

    # date
    if 'date' in kie_dict:
        kie_dict['date'] = postprocess_date(kie_dict['date'])

    #### normalize number
    number_fields = ['total_value', 'VAT_amount_value']
    for number_field in number_fields:
        if number_field not in kie_dict:
            continue
        kie_dict[number_field] = [normalize_number_wordgroup(word_group) for word_group in kie_dict[number_field]]

    #tax code value 
    tax_code_fields = ['buyer_tax_code_value', 'seller_tax_code_value']
    for tax_code_field in tax_code_fields:
        if tax_code_field in kie_dict:

            kie_dict[tax_code_field] = [normalize_tax_wordgroup(word_group, reserve_minus=True) for word_group in kie_dict[tax_code_field]]
        
    # mobile
    if "seller_mobile_value" in kie_dict:
        kie_dict['seller_mobile_value'] = [normalize_number_wordgroup(word_group, reserve_plus=True) for word_group in kie_dict['seller_mobile_value']]
        
    # if 'total_in_words_value' in kie_dict:
    #     kie_dict['total_in_words_value'] = postprocess_total_in_words(kie_dict['total_in_words_value'])



    return kie_dict