"""Utils for the VisualWebBench benchmark. Referenced from https://github.com/VisualWebBench/VisualWebBench/blob/main/utils/eval_utils.py"""

import re

import numpy as np
import torch
from rouge import Rouge
from torchvision.ops import box_iou


def rouge_score(preds: list[str], golds: list[str], **kwargs) -> float:
    rouge = Rouge(metrics=["rouge-1", "rouge-2", "rouge-l"])
    scores: dict[str, dict[str, float]] = rouge.get_scores(preds, golds, avg=True)  # type: ignore

    rouge1_score = scores["rouge-1"]["f"]
    rouge2_score = scores["rouge-2"]["f"]
    rouge_l_score = scores["rouge-l"]["f"]

    rouge_avg = float(np.mean([rouge1_score, rouge2_score, rouge_l_score]))

    # To get the retrieve the metrics name from the function
    # return dict(rouge_avg=rouge_avg)
    return rouge_avg


def eval_heading_ocr_or_web_caption(preds: list[str], golds: list[str], **kwargs) -> float:
    assert len(preds) == len(golds)
    for i in range(len(preds)):
        if not preds[i]:
            preds[i] = " "
        if not golds[i]:
            golds[i] = " "

    return rouge_score(preds, golds)


def eval_element_ocr(preds: list[str], golds: list[str], **kwargs) -> float:
    assert len(preds) == len(golds)
    for i in range(len(preds)):
        if not preds[i] or len(preds[i]) == 1:
            preds[i] = " "
        if not golds[i]:
            golds[i] = " "

    return rouge_score(preds, golds)


def eval_element_or_action(preds: list[str], golds: list[int], **kwargs) -> float:
    results: list[bool] = []

    for pred, gold in zip(preds, golds):
        str_pred = parse_multi_choice_response(pred, [chr(ord("A") + i) for i in range(8)])
        try:
            if ord("A") <= ord(str_pred) <= ord("Z"):
                idx_pred = ord(str_pred) - ord("A")
            else:
                idx_pred = -1
        except Exception:
            idx_pred = -1
        results.append(idx_pred == gold)

    # To get the retrieve the metrics name from the function
    # return dict(accuracy=sum(results) / len(results) * 100)
    return sum(results) / len(results)


def eval_element_bbox_ground(
    preds: list[tuple[float, float, float, float]], golds: list[tuple[float, float, float, float]], **kwargs
) -> float:
    # print('preds[0]', preds[0])
    # print('golds[0]', golds[0])
    correct = total_cnt = 0
    for i, predict_bbox in enumerate(preds):
        if not predict_bbox:
            predict_bbox = (0.0, 0.0, 0.0, 0.0)
        try:
            target_bbox = torch.tensor(golds[i], dtype=torch.float32).view(-1, 4)
            tensor_predict_bbox = torch.tensor(predict_bbox, dtype=torch.float32).view(-1, 4)
            iou = box_iou(tensor_predict_bbox, target_bbox)
            iou = iou.item()
            if iou >= 0.5:
                correct += 1
        except Exception:
            pass

        total_cnt += 1

    # To get the retrieve the metrics name from the function
    # return dict(precision=correct / total_cnt * 100)
    return correct / total_cnt


def eval_action_bbox_ground(
    preds: list[tuple[float, float, float, float]], golds: list[tuple[float, float, float, float]], **kwargs
) -> float:
    correct = total_cnt = 0
    for i, predict_bbox in enumerate(preds):
        if not predict_bbox:
            predict_bbox = (0.0, 0.0, 0.0, 0.0)
        try:
            target_bbox = torch.tensor(golds[i], dtype=torch.float32).view(-1, 4)
            tensor_predict_bbox = torch.tensor(predict_bbox, dtype=torch.float32).view(-1, 4)
            iou = box_iou(tensor_predict_bbox, target_bbox)
            iou = iou.item()
            if iou >= 0.5:
                correct += 1
        except Exception:
            pass

        total_cnt += 1

    # To get the retrieve the metrics name from the function
    # return dict(precision=correct / total_cnt * 100)
    return correct / total_cnt


def eval_webqa(preds: list[str], golds: list[list[str]], **kwargs) -> float:
    f1_scores = []
    rouge = Rouge(metrics=["rouge-1"])
    for pred, gold_list in zip(preds, golds):
        try:
            if not pred:
                pred = " "
            cur_f1 = max([rouge.get_scores([pred], [gold], avg=True)["rouge-1"]["f"] for gold in gold_list])  # type: ignore
            f1_scores.append(cur_f1)
        except Exception:
            pass

    # To get the retrieve the metrics name from the function
    # return dict(f1=sum(f1_scores) / len(f1_scores) * 100)
    return sum(f1_scores) / len(f1_scores)


def eval_element_point_ground(
    preds: list[tuple[float, float]], golds: list[tuple[float, float, float, float]], **kwargs
) -> float:
    acc_lst = []
    for pred, gold in zip(preds, golds):
        x, y = pred
        left, top, right, bottom = gold
        acc_lst.append(left <= x <= right and top <= y <= bottom)
    # To get the retrieve the metrics name from the function
    # return dict(accuracy=sum(acc_lst) / len(acc_lst) * 100)
    return sum(acc_lst) / len(acc_lst)


def eval_action_point_ground(
    preds: list[tuple[float, float]], golds: list[tuple[float, float, float, float]], **kwargs
) -> float:
    acc_lst = []
    for pred, gold in zip(preds, golds):
        x, y = pred
        left, top, right, bottom = gold
        acc_lst.append(left <= x <= right and top <= y <= bottom)
    # To get the retrieve the metrics name from the function
    # return dict(accuracy=sum(acc_lst) / len(acc_lst) * 100)
    return sum(acc_lst) / len(acc_lst)


# ----------- Process Multi-choice -------------
def parse_multi_choice_response(response: str, all_choices: list[str]) -> str:
    """
    Parse the prediction from the generated response.
    Return the predicted index e.g., A, B, C, D.
    """
    if len(response) == 1:
        return response.upper()
    elif not response:
        return "a"
    elif re.match(r"[A-Z]\.", response):
        return response[0]

    for char in [",", ".", "!", "?", ";", ":", "'", '"']:
        response = response.replace(char, "")
    response = " " + response + " "  # add space to avoid partial match

    ans_with_brack = False
    candidates = []
    for choice in all_choices:  # e.g., (A) (B) (C) (D)
        if f"({choice})" in response:
            candidates.append(choice)
            ans_with_brack = True

    if len(candidates) == 0:
        for choice in all_choices:  # e.g., A B C D
            if f" {choice} " in response:
                candidates.append(choice)

    if len(candidates) == 0:  # still not get answer
        # pred_index = random.choice(all_choices)
        pred_index = "z"
    elif len(candidates) > 1:
        start_indexes = []
        if ans_with_brack:
            for can in candidates:
                index = response.rfind(f"({can})")
                start_indexes.append(index)  # -1 will be ignored anyway
            # start_indexes = [generated_response.index(f'({can})') for can in candidates]
        else:
            for can in candidates:
                index = response.rfind(f" {can} ")
                start_indexes.append(index)
        # get the last one
        pred_index = candidates[np.argmax(start_indexes)]
    else:  # if only one candidate, use it.
        pred_index = candidates[0]

    return pred_index
