import cv2
import numpy as np
from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment


class SimpleAffineTransform:
    def __init__(self, translation=(0, 0), scale=1.0):
        self.translation = np.array(translation)
        self.scale = scale

    def estimate(self, src, dst):
        src_center = np.mean(src, axis=0)
        dst_center = np.mean(dst, axis=0)
        self.translation = dst_center - src_center
        src_dists = np.linalg.norm(src - src_center, axis=1)
        dst_dists = np.linalg.norm(dst - dst_center, axis=1)
        self.scale = np.mean(dst_dists) / (np.mean(src_dists) + 1e-10)

    def inverse(self):
        return SimpleAffineTransform(-self.translation, 1.0 / self.scale)

    def __call__(self, coords):
        return (
            self.scale * (coords - np.mean(coords, axis=0))
            + np.mean(coords, axis=0)
            + self.translation
        )

    def residuals(self, src, dst):
        return np.sqrt(np.sum((self(src) - dst) ** 2, axis=1))


def norm_coords(x, left, right):
    if x < left:
        return left
    if x > right:
        return right
    return x


def norm_same_token(token):
    special_map = {
        "\\dot": ".",
        "\\Dot": ".",
        "\\cdot": ".",
        "\\cdotp": ".",
        "\\ldotp": ".",
        "\\mid": "|",
        "\\rightarrow": "\\to",
        "\\top": "T",
        "\\Tilde": "\\tilde",
        "\\prime": "'",
        "\\ast": "*",
        "\\left<": "\\langle",
        "\\right>": "\\rangle",
        "\\lbrace": "\{",
        "\\rbrace": "\}",
        "\\lbrack": "[",
        "\\rbrack": "]",
        "\\blackslash": "/",
        "\\slash": "/",
        "\\leq": "\\le",
        "\\geq": "\\ge",
        "\\neq": "\\ne",
        "\\Vert": "\\|",
        "\\lVert": "\\|",
        "\\rVert": "\\|",
        "\\vert": "|",
        "\\lvert": "|",
        "\\rvert": "|",
        "\\colon": ":",
        "\\Ddot": "\\ddot",
        "\\Bar": "\\bar",
        "\\Vec": "\\vec",
        "\\parallel": "\\|",
        "\\dag": "\\dagger",
        "\\ddag": "\\ddagger",
        "\\textlangle": "<",
        "\\textrangle": ">",
        "\\textgreater": ">",
        "\\textless": "<",
        "\\textbackslash": "\\",
        "\\textunderscore": "_",
        "\\=": "=",
        "\\neg": "\\lnot",
        "\\neq": "\\not=",
    }
    if token.startswith("\\left") or token.startswith("\\right"):
        if (
            "arrow" not in token
            and "<" not in token
            and ">" not in token
            and "harpoon" not in token
        ):
            token = token.replace("\\left", "").replace("\\right", "")
    if token.startswith("\\big") or token.startswith("\\Big"):
        if "\\" in token[4:]:
            token = "\\" + token[4:].split("\\")[-1]
        else:
            token = token[-1]
    if token in special_map.keys():
        token = special_map[token]
    if token.startswith("\\wide"):
        return token.replace("wide", "")
    if token.startswith("\\var"):
        return token.replace("var", "")
    if token.startswith("\\string"):
        return token.replace("\\string", "")
    return token


class HungarianMatcher:
    def __init__(
        self,
        cost_token: float = 1,
        cost_position: float = 0.05,
        cost_order: float = 0.15,
    ):
        self.cost_token = cost_token
        self.cost_position = cost_position
        self.cost_order = cost_order
        self.cost = {}

    def calculate_token_cost(self, box_gt, box_pred):

        all_tokens = [data["token"] for data in box_gt + box_pred]
        unique_tokens = sorted(list(set(all_tokens)))
        token2id = {token: i for i, token in enumerate(unique_tokens)}
        num_classes = len(token2id)

        all_norm_tokens = [norm_same_token(data["token"]) for data in box_gt + box_pred]
        unique_norm_tokens = sorted(list(set(all_norm_tokens)))
        token2id_norm = {token: i for i, token in enumerate(unique_norm_tokens)}
        num_classes_norm = len(token2id_norm)

        gt_token_array = np.array([token2id[data["token"]] for data in box_gt])
        norm_gt_token_array = np.array(
            [token2id_norm[norm_same_token(data["token"])] for data in box_gt]
        )

        pred_token_logits = np.zeros((len(box_pred), num_classes))
        for i, data in enumerate(box_pred):
            if data["token"] in token2id:
                pred_token_logits[i, token2id[data["token"]]] = 1

        norm_pred_token_logits = np.zeros((len(box_pred), num_classes_norm))
        for i, data in enumerate(box_pred):
            norm_token = norm_same_token(data["token"])
            if norm_token in token2id_norm:
                norm_pred_token_logits[i, token2id_norm[norm_token]] = 1

        if gt_token_array.size == 0 or pred_token_logits.shape[0] == 0:
            return np.empty((len(box_gt), len(box_pred)))

        token_cost = 1.0 - pred_token_logits[:, gt_token_array]
        norm_token_cost = 1.0 - norm_pred_token_logits[:, norm_gt_token_array]

        token_cost[np.logical_and(token_cost == 1, norm_token_cost == 0)] = 0.005
        return token_cost.T

    def box2array(self, box_list, size):
        W, H = size
        box_array = []
        for box in box_list:
            x_min, y_min, x_max, y_max = box["bbox"]
            box_array.append(
                [
                    (x_min + x_max) / (2 * W),
                    (y_min + y_max) / (2 * H),
                    (x_max - x_min) / W,
                    (y_max - y_min) / H,
                ]
            )
        return np.array(box_array)

    def order2array(self, box_list, max_token_lens=None):
        if not max_token_lens:
            max_token_lens = len(box_list)
        return np.array([[idx / max_token_lens] for idx, _ in enumerate(box_list)])

    def calculate_l1_cost(self, gt_array, pred_array):
        if gt_array.shape[0] == 0 or pred_array.shape[0] == 0:
            return np.empty((gt_array.shape[0], pred_array.shape[0]))
        return cdist(gt_array, pred_array, "minkowski", p=1) / gt_array.shape[-1]

    def __call__(self, box_gt, box_pred, gt_size, pred_size):
        if not box_gt or not box_pred:
            return []
        gt_box_array = self.box2array(box_gt, gt_size)
        pred_box_array = self.box2array(box_pred, pred_size)
        max_token_lens = max(len(box_gt), len(box_pred))
        gt_order_array = self.order2array(box_gt, max_token_lens)
        pred_order_array = self.order2array(box_pred, max_token_lens)
        token_cost = self.calculate_token_cost(box_gt, box_pred)
        position_cost = self.calculate_l1_cost(gt_box_array, pred_box_array)
        order_cost = self.calculate_l1_cost(gt_order_array, pred_order_array)
        self.cost = {
            "token": token_cost,
            "position": position_cost,
            "order": order_cost,
        }
        cost = (
            self.cost_token * token_cost
            + self.cost_position * position_cost
            + self.cost_order * order_cost
        )
        cost[np.isnan(cost) | np.isinf(cost)] = 100
        row_ind, col_ind = linear_sum_assignment(cost)
        return list(zip(row_ind, col_ind))


def update_inliers(ori_inliers, sub_inliers):
    inliers = np.copy(ori_inliers)
    sub_idx = -1
    for idx in range(len(ori_inliers)):
        if ori_inliers[idx] == False:
            sub_idx += 1
            if sub_inliers[sub_idx] == True:
                inliers[idx] = True
    return inliers