# Copyright 2025 Google LLC
#
# 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
#
#     https://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.

%pip install --upgrade --user --quiet google-cloud-aiplatform google-genai

import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

# Use the environment variable if the user doesn't provide Project ID.
import os

import vertexai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on

if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

vertexai.init(project=PROJECT_ID, location=LOCATION)

import abc
import dataclasses
import random
import re
from collections.abc import Callable

import pandas as pd
from google import genai
from scipy.stats import kendalltau, spearmanr
from sklearn.metrics import cohen_kappa_score, confusion_matrix

autorater_model = "gemini-2.5-flash"  # @param {type:"string", isTemplate: true}

@dataclasses.dataclass
class Example:
    """Represents a single example for rating.

    Attributes:
        prompt: The prompt string.
        response1: The first response string.
        response2: The second response string.
    """

    prompt: str
    response1: str
    response2: str


@dataclasses.dataclass
class RaterInstruction:
    """Stores a rating prompt and its associated result parser.

    Attributes:
        prompt: The prompt string for the rater.
        result_parser: A function to parse the rater's output string into an integer score.
    """

    prompt: str
    result_parser: Callable[[str], int | None]


def simple_no_tie_result_parser(result_str: str) -> int | None:
    """Parses a result string for simple prompts (no ties allowed).

    Args:
        result_str: The rater's output string.

    Returns:
        -1 if response 1 is better, 1 if response 2 is better, or None if the output is invalid.
    """
    matches = re.findall(r"<winner>(.*?)</winner>", result_str)
    if not matches or len(matches) > 1:
        return None
    if matches[0] == "1":
        return -1
    if matches[0] == "2":
        return 1
    return None


def simple_with_tie_result_parser(result_str: str) -> int | None:
    """Parses a result string for simple prompts (ties allowed).

    Args:
        result_str: The rater's output string.

    Returns:
        -1 if response 1 is better, 1 if response 2 is better, 0 if it's a tie, or None if the output is invalid.
    """
    if "<tie>" in result_str:
        return 0

    matches = re.findall(r"<winner>(.*?)</winner>", result_str)
    if not matches or len(matches) > 1:
        return None
    if matches[0] == "1":
        return -1
    if matches[0] == "2":
        return 1
    return None


def reward_bench_gemini_no_tie_parser(result_str: str) -> int | None:
    """Parses a result string for reward-bench Gemini prompts (no ties).

    Args:
        result_str: The rater's output string.

    Returns:
        -1 if response 1 (A) is better, 1 if response 2 (B) is better, or None if the output is invalid.
    """
    match = re.search(r"\[\[(A|B)\]\]", result_str)
    if match:
        if match.group(1) == "A":
            return -1
        if match.group(1) == "B":
            return 1
        return None
    return None


PROMPT_MAP = {
    "simple_no_tie": RaterInstruction(
        prompt="""
          ## You are an impartial judge to evaluate the quality of two responses to the user prompt.

          [Start of Prompt]{prompt} [end of  prompt]
          [Start of Response 1]{response1}[End of the response 1]
          [Start of Response 2]{response2}[End of the response 2]

          Your output should strictly follow this format: <winner>1</winner>, if response 1 is better; <winner>2</winner>, if response 2 is better.
          """,
        result_parser=simple_no_tie_result_parser,
    ),
    "simple_with_tie": RaterInstruction(
        prompt="""
          ## You are an impartial judge to evaluate the quality of two responses to the user prompt.

          [Start of Prompt]{prompt} [end of  prompt]
          [Start of Response 1]{response1}[End of the response 1]
          [Start of Response 2]{response2}[End of the response 2]

          Your output should strictly follow this format: <winner>1</winner>, if response 1 is better; <winner>2</winner>, if response 2 is better; <tie> if you cannot determine a winner
          """,
        result_parser=simple_with_tie_result_parser,
    ),
    "reward-bench_gemini_no_tie": RaterInstruction(
        prompt=(
            "Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user question displayed below. "
            "You should choose the assistant that follows the user's instructions and answers the user's question better. "
            "Your evaluation should consider factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of their responses. "
            "Avoid any position biases and ensure that the order in which the responses were presented does not influence your decision. "
            "Do not allow the length of the responses to influence your evaluation. Do not favor certain names of the assistants. "
            "Be as objective as possible. "
            "Your output should only consist of '[[A]]' if assistant A is better, or '[[B]]' if assistant B is better. Omit any other output.\n"
            "[User Question]\n{prompt}\n\n[The Start of Assistant A's Answer]\n{response1}\n[The End of Assistant A's Answer]\n\n[The Start of Assistant B's Answer]\n{response2}\n[The End of Assistant B's Answer]",
        ),
        result_parser=simple_no_tie_result_parser,
    ),
}


class AutoRater(abc.ABC):
    """Abstract base class for automatic rating of responses."""

    def __init__(self, name, model_id, prompt_method):
        """Initializes the AutoRater.

        Args:
        name: The name of the rater.
        model_id: The model id to use for rating.
        prompt_method: The name of the prompt method to use.
        """
        self.name = name
        self.model_id = model_id
        self.prompt_method = prompt_method
        self.client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

    @abc.abstractmethod
    def rate_one_example(self, example: Example) -> int:
        """Rates a single example.

        Args:
        example: The example to rate.

        Returns:
        The rating for the example.
        """
        # for pair-wise: -1, 0, 1
        # for point-wise: scale 1-5
        # This method should also be implemented by subclasses

    def rate_batch(self, id_to_example: dict[str, Example]) -> dict[str, int]:
        """Rates a batch of examples.

        Args:
        id_to_example: A dictionary mapping example IDs to examples.

        Returns:
        A dictionary mapping example IDs to ratings.
        """
        id_to_rating = {}
        for id, ex in id_to_example.items():
            rating = self.rate_one_example(ex)
            if rating is not None:
                id_to_rating[id] = rating
        return id_to_rating


class BasicRater(AutoRater):
    """A basic rater that uses a single call to the language model."""

    def __init__(self, model_id, prompt_method):
        """Initializes the BasicRater.

        Args:
        model_id: The model id.
        prompt_method: The prompt method name.
        """
        super().__init__("Basic", model_id, prompt_method)

    def rate_one_example(self, example: Example) -> int | None:
        """Rates a single example using a single LLM call.

        Args:
        example: The example to rate.

        Returns:
        The rating for the example, or None if an error occurred.
        """
        autorater_prompt = PROMPT_MAP[self.prompt_method].prompt.format(
            prompt=example.prompt,
            response1=example.response1,
            response2=example.response2,
        )
        try:
            response = self.client.models.generate_content(
                model=self.model_id, contents=autorater_prompt
            )
        except Exception as e:
            print(e)
            return None
        if response.candidates and response.candidates[0].content.parts:
            result_str = response.candidates[0].content.parts[0].text
            print(result_str)
            result = PROMPT_MAP[self.prompt_method].result_parser(result_str)
            print(result)
            return result
        return None


class SelfConsistencyRater(AutoRater):
    """A self-consistency rater that uses multiple calls to the language model and aggregates the results."""

    def __init__(self, model_id, prompt_method, num_calls):
        """Initializes the SelfConsistencyRater.

        Args:
        model_id: The model id.
        prompt_method: The prompt method name.
        num_calls: The number of calls to make to the language model.
        """
        super().__init__("SelfConsistency", model_id, prompt_method)
        self.num_calls = num_calls

    def _get_consensus_result(self, result_list: list[int]) -> int:
        """Aggregates the results from multiple calls to the language model.

        Args:
        result_list: A list of integer results.

        Returns:
        The aggregated result.
        """
        avg = sum(result_list) / len(result_list)
        if avg > 0.5:
            return 1
        if avg < -0.5:
            return -1
        return 0

    def rate_one_example(self, example: Example) -> int | None:
        """Rates a single example using multiple LLM calls and consensus aggregation.

        Args:
        example: The example to rate.


        Returns:
        The aggregated rating for the example, or None if an error occurred or no valid results were obtained.
        """
        result_list = []
        for i in range(self.num_calls):
            autorater_prompt = PROMPT_MAP[self.prompt_method].prompt.format(
                prompt=example.prompt,
                response1=example.response1,
                response2=example.response2,
            )
            try:
                response = self.client.models.generate_content(
                    model=self.model_id, contents=autorater_prompt
                )
            except Exception as e:
                print(e)
                continue

            if response.candidates and response.candidates[0].content.parts:
                result_str = response.candidates[0].content.parts[0].text
                print(result_str)
                result = PROMPT_MAP[self.prompt_method].result_parser(result_str)
                print(result)
                if result is not None:
                    result_list.append(result)
        print(result_list)
        if result_list:
            result = self._get_consensus_result(result_list)
            print(result)
            return result
        return None

def process_reward_bench(split: str = "filtered", total_count: int = None):
    """Processes reward benchmark data.

    This function processes reward benchmark data based on the specified split
    and total count. If 'total_count' is less than 0, all data will be used.

    Args:
      split: The split of data to use ('filtered' by default).
      total_count: The maximum number of data points to process. If not set,
        all data will be used.

    Returns:
      This function does not return any value (implicitly returns None).
    """
    split_to_path = {
        "raw": "data/raw-00000-of-00001.parquet",
        "filtered": "data/filtered-00000-of-00001.parquet",
    }
    df = pd.read_parquet("hf://datasets/allenai/reward-bench/" + split_to_path[split])
    id_to_example = {}
    id_to_golden_rating = {}

    for index, (id, row) in enumerate(df.iterrows()):
        if total_count is not None and index > total_count - 1:
            break
        question_id = f"{row['subset']}-{row['id']}"
        golden_rating = random.choice([-1, 1])
        if golden_rating == -1:
            response1 = row["chosen"]
            response2 = row["rejected"]
        else:
            response2 = row["chosen"]
            response1 = row["rejected"]
        prompt = row["prompt"]
        example = Example(prompt=prompt, response1=response1, response2=response2)
        id_to_example[question_id] = example
        id_to_golden_rating[question_id] = golden_rating

    return id_to_example, id_to_golden_rating

def get_caliboration_result(
    model_outputs: list[int],
    golden_labels: list[int],
    labels: list[int] = [-1, 0, 1],
    weights: str = "quadratic",
):
    """Calculates various evaluation metrics for model outputs against golden labels.

    Args:
        model_outputs: Predicted labels from the model.
        golden_labels: True labels.
        labels: Unique labels in the dataset. Default is [-1, 0, 1].
        weights: Weighting scheme for Cohen's Kappa. Default is "quadratic".

    Returns:
        A tuple containing the confusion matrix (DataFrame), Cohen's Kappa score, Spearman's rank correlation, and Kendall's rank correlation.
    """
    spearman, _ = spearmanr(model_outputs, golden_labels)
    kendall, _ = kendalltau(model_outputs, golden_labels)
    kappa = cohen_kappa_score(
        model_outputs, golden_labels, labels=labels, weights=weights
    )

    conf_matrix = confusion_matrix(golden_labels, model_outputs, labels=labels)
    conf_matrix_df = pd.DataFrame(
        conf_matrix,
        index=["Gold_1", "Gold_Tie", "Gold_2"],
        columns=["Model_1", "Model_Tie", "Model_2"],
    )

    return conf_matrix_df, kappa, spearman, kendall


def get_aligned_lists(
    id_to_auto_ratings: dict, id_to_golden_rating: dict
) -> tuple[list, list]:
    """Aligns two dictionaries of ratings based on shared keys (IDs).

    Args:
        id_to_auto_ratings: Dictionary mapping IDs to automatically generated ratings.
        id_to_golden_rating: Dictionary mapping IDs to golden ratings.

    Returns:
        A tuple containing two aligned lists of ratings.
    """
    id_list = sorted(id_to_auto_ratings.keys())
    auto_ratings = [id_to_auto_ratings[id] for id in id_list]
    golden_ratings = [id_to_golden_rating[id] for id in id_list]
    return auto_ratings, golden_ratings

id_to_example, id_to_golden_rating = process_reward_bench(
    split="filtered", total_count=100
)
len(id_to_example)

rater = BasicRater(autorater_model, "simple_no_tie")
id_to_basic_auto_ratings = rater.rate_batch(id_to_example)

auto_ratings, golden_ratings = get_aligned_lists(
    id_to_basic_auto_ratings, id_to_golden_rating
)
conf_matrix_df, kappa_score, corr_spearman, corr_kendall = get_caliboration_result(
    model_outputs=auto_ratings,
    golden_labels=golden_ratings,
    labels=[-1, 0, 1],
    weights="quadratic",
)
print(f"kappa_score: {kappa_score}")
print(f"corr_spearman: {corr_spearman}")
print(f"corr_kendall: {corr_kendall}")

print(conf_matrix_df)

sc_rater = SelfConsistencyRater(autorater_model, "simple_no_tie", 3)
id_to_sc_auto_ratings = sc_rater.rate_batch(id_to_example)

auto_ratings, golden_ratings = get_aligned_lists(
    id_to_sc_auto_ratings, id_to_golden_rating
)
print(len(auto_ratings))
print(len(golden_ratings))

conf_matrix_df, kappa_score, corr_spearman, corr_kendall = get_caliboration_result(
    model_outputs=auto_ratings,
    golden_labels=golden_ratings,
    labels=[-1, 0, 1],
    weights="quadratic",
)
print(f"kappa_score: {kappa_score}")
print(f"corr_spearman: {corr_spearman}")
print(f"corr_kendall: {corr_kendall}")

print(conf_matrix_df)