# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
from functools import partial
from itertools import combinations
from typing import Any, Dict, List

import datasets
import numpy as np

from lm_eval.tasks.score import utils
from lm_eval.tasks.score.math.math_grader import (
    extract_answer,
    math_equal,
    normalize_answer_string,
)
from lm_eval.tasks.score.utils import robustness_doc_to_text
from lm_eval.utils import eval_logger


TEMPLATE_FILE_PATH = os.path.join(os.path.dirname(__file__), "prompt_templates.json")

PROMPT_ROBUSTNESS_TEMPLATE_KEY = "prompt_robustness"
NON_GREEDY_ROBUSTNESS_TEMPLATE_KEY = "non_greedy_robustness"

math_robustness_doc_to_text = robustness_doc_to_text


def find_boxed_entries(answer_str):
    stack = []
    results = []
    i = 0

    while i < len(answer_str):
        if answer_str[i : i + 7] == "\\boxed{":
            stack.append(i + 7)
            i += 7
        elif answer_str[i] == "{":
            if stack:
                stack.append(i + 1)
            i += 1
        elif answer_str[i] == "}":
            if stack:
                start = stack.pop()
                if not stack:
                    results.append(answer_str[start:i])
            i += 1
        else:
            i += 1

    if len(results) == 0:
        raise ValueError("Not enough boxed entries")
    else:
        results = [normalize_answer_string(result) for result in results]

    if len(results) == 1:
        # Single boxed entry, trivial case
        return results

    else:
        # Multiple boxed entries. There are two cases possible
        # (a) The reference solution has the same question answered in multiple ways
        # (b) The answer is split across multiple boxed entries and we need to merge
        result_equal = True
        for idx in range(len(results) - 1):
            if not (results[idx] == results[idx + 1]):
                result_equal = False
                break

        if result_equal:
            # Same problem solved in multiple ways
            return [results[0]]
        else:
            return results


def extract_answer_dataset(solution: str, problem: str, corrected_answers: list) -> str:
    entries = find_boxed_entries(solution)

    if len(entries) == 1:
        parsed_answer = entries[0]

    if len(entries) > 1:
        for item in corrected_answers:
            if item["problem"] == problem:
                parsed_answer = item["answer"]
                break
        else:
            parsed_answer = ", ".join(entries)

    if not (
        ("Find the equation" in problem)
        or ("Enter the equation" in problem)
        or ("What is the equation" in problem)
        or ("described by the equation" in problem)
        or ("Find an equation" in problem)
    ) and ("=" in parsed_answer):
        if parsed_answer.count("=") == 1:
            # For greater count, it means we're just predicting values of multiple variables
            parsed_answer = parsed_answer.split("=")[1]
    return parsed_answer


def process_docs(dataset: datasets.Dataset) -> datasets.Dataset:
    def _process_doc(doc: dict, idx, corrected_answer) -> dict:
        out_doc = {
            "question": doc["problem"],
            "question_id": idx,
            "solution": doc["solution"],
            "answer": extract_answer_dataset(
                doc["solution"], doc["problem"], corrected_answer
            ),
        }
        return out_doc

    corrected_answer_path = os.path.join(
        os.path.dirname(__file__), "to_be_fixed_questions.json"
    )

    with open(corrected_answer_path, "r") as f:
        corrected_answers = json.load(f)

    return dataset.map(
        partial(_process_doc, corrected_answer=corrected_answers), with_indices=True
    )


def prompt_robustness_process_docs(doc: datasets.Dataset) -> datasets.Dataset:
    doc = process_docs(doc)
    return utils.process_docs_add_prompts(
        doc,
        templates_key=PROMPT_ROBUSTNESS_TEMPLATE_KEY,
        template_file_path=TEMPLATE_FILE_PATH,
    )


def non_greedy_robustness_process_docs(doc: datasets.Dataset) -> datasets.Dataset:
    doc = process_docs(doc)
    return utils.non_greedy_robustness_process_docs(
        doc,
        templates_key=NON_GREEDY_ROBUSTNESS_TEMPLATE_KEY,
        template_file_path=TEMPLATE_FILE_PATH,
    )


def process_results(doc: dict, results: List[str]) -> Dict[str, int]:
    answer = extract_answer(results[0])

    if math_equal(answer, doc["answer"]):
        retval = 1
    else:
        retval = 0

    prompt_id = doc["prompt_id"]

    results = {
        f"{prompt_id}_accuracy": (prompt_id, retval),
        "consistency_rate": (doc["question_id"], answer),
    }
    return results


def non_greedy_robustness_process_results(
    doc: dict, results: List[str]
) -> Dict[str, int]:
    answer = extract_answer(results[0])
    return {"non_greedy_accuracy": (doc["question_id"], answer, doc["answer"], None)}


def per_prompt_accuracy(results: List[Dict[str, Any]], p_id=0) -> float:
    accuracies = []
    for result in results:
        prompt_id, retval = result
        if prompt_id != p_id:
            continue
        accuracies.append(retval)

    accuracy = sum(accuracies) / len(accuracies)
    eval_logger.info(f"Prompt - {prompt_id} accuracy: {accuracy}")

    return np.round(accuracy, 4)


per_prompt_accuracy_0 = partial(per_prompt_accuracy, p_id=0)
per_prompt_accuracy_1 = partial(per_prompt_accuracy, p_id=1)
per_prompt_accuracy_2 = partial(per_prompt_accuracy, p_id=2)
per_prompt_accuracy_3 = partial(per_prompt_accuracy, p_id=3)
per_prompt_accuracy_4 = partial(per_prompt_accuracy, p_id=4)
per_prompt_accuracy_5 = partial(per_prompt_accuracy, p_id=5)
per_prompt_accuracy_6 = partial(per_prompt_accuracy, p_id=6)
per_prompt_accuracy_7 = partial(per_prompt_accuracy, p_id=7)
per_prompt_accuracy_8 = partial(per_prompt_accuracy, p_id=8)
per_prompt_accuracy_9 = partial(per_prompt_accuracy, p_id=9)


def calculate_consistency_rate(responses: List[List[str]]) -> float:
    """
    Calculate the Consistency Rate (CR) for a given set of responses.

    Args:
    responses: List of lists, where each inner list contains responses to the same question.

    Returns:
    The consistency rate as a float.
    """
    total_similarity = 0
    total_combinations = 0

    for response_set in responses:
        pairs = combinations(response_set, 2)
        num_pairs = len(response_set) * (len(response_set) - 1) / 2
        total_combinations += num_pairs
        for answer1, answer2 in pairs:
            total_similarity += int(math_equal(answer1, answer2))

    return total_similarity / total_combinations if total_combinations > 0 else 0.0


def math_prompt_consistency_rate(results: List[Dict[str, Any]]) -> float:
    """
    Calculate the Consistency Rate (CR) for a given set of responses.

    Args:
    responses: List of lists, where each inner list contains responses to the same question.

    Returns:
    The consistency rate as a float.
    """
    question_answers_dict = {}

    for result in results:
        question_id, answer = result
        if question_id not in question_answers_dict:
            question_answers_dict[question_id] = []
        question_answers_dict[question_id].append(answer)

    question_answers_list = [answers for answers in question_answers_dict.values()]

    return calculate_consistency_rate(question_answers_list)


def non_greedy_accuracy(results: List[Dict[str, Any]]) -> float:
    accuracies = []
    for result in results:
        question_id, final_answer, gt, _ = result
        if math_equal(final_answer, gt):
            retval = 1
        else:
            retval = 0
        accuracies.append(retval)

    accuracy = sum(accuracies) / len(accuracies)
    eval_logger.info(f"Non greedy accuracy: {accuracy}")

    return np.round(accuracy, 4)
