# Copyright 2024 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[evaluation]

# 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()

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
LOCATION = "us-central1"  # @param {type:"string"}
EXPERIMENT = "eval-migration-ga"  # @param {type:"string"}

if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    raise ValueError("Please set your PROJECT_ID")

import vertexai

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

import pandas as pd
from vertexai.evaluation import (
    EvalTask,
    MetricPromptTemplateExamples,
    PairwiseMetric,
    PointwiseMetric,
)
from vertexai.preview.evaluation import notebook_utils

questions = [
    "Which part of the brain does short-term memory seem to rely on?",
    "What provided the Roman senate with exuberance?",
    "What area did the Hasan-jalalians command?",
]

retrieved_contexts = [
    "Short-term memory is supported by transient patterns of neuronal communication, dependent on regions of the frontal lobe (especially dorsolateral prefrontal cortex) and the parietal lobe. Long-term memory, on the other hand, is maintained by more stable and permanent changes in neural connections widely spread throughout the brain. The hippocampus is essential (for learning new information) to the consolidation of information from short-term to long-term memory, although it does not seem to store information itself. Without the hippocampus, new memories are unable to be stored into long-term memory, as learned from patient Henry Molaison after removal of both his hippocampi, and there will be a very short attention span. Furthermore, it may be involved in changing neural connections for a period of three months or more after the initial learning.",
    "In 62 BC, Pompey returned victorious from Asia. The Senate, elated by its successes against Catiline, refused to ratify the arrangements that Pompey had made. Pompey, in effect, became powerless. Thus, when Julius Caesar returned from a governorship in Spain in 61 BC, he found it easy to make an arrangement with Pompey. Caesar and Pompey, along with Crassus, established a private agreement, now known as the First Triumvirate. Under the agreement, Pompey's arrangements would be ratified. Caesar would be elected consul in 59 BC, and would then serve as governor of Gaul for five years. Crassus was promised a future consulship.",
    "The Seljuk Empire soon started to collapse. In the early 12th century, Armenian princes of the Zakarid noble family drove out the Seljuk Turks and established a semi-independent Armenian principality in Northern and Eastern Armenia, known as Zakarid Armenia, which lasted under the patronage of the Georgian Kingdom. The noble family of Orbelians shared control with the Zakarids in various parts of the country, especially in Syunik and Vayots Dzor, while the Armenian family of Hasan-Jalalians controlled provinces of Artsakh and Utik as the Kingdom of Artsakh.",
]

generated_answers = [
    "frontal lobe and the parietal lobe",
    "The Roman Senate was filled with exuberance due to successes against Catiline.",
    "The Hasan-Jalalians commanded the area of Syunik and Vayots Dzor.",
]

baseline_answers = [
    "the frontal cortex and the parietal cortex, which are crucial for sensory and cognitive functions",
    "The Roman Senate celebrated triumphantly after significant victories over Catiline, bolstering their political influence",
    "The Hasan-Jalalians held control over the regions of Syunik and Vayots Dzor, maintaining power through strategic alliances and military strength",
]

eval_dataset = pd.DataFrame(
    {
        "instruction": questions,
        "context": retrieved_contexts,
        "response": generated_answers,
    }
)

relevance_prompt_template = """
# Instruction
You are a professional writing evaluator. Your job is to score writing responses according to pre-defined evaluation criteria.
You will be assessing question answering relevance, which measures the ability to respond with relevant information when asked a question.
You will assign the writing response a score from 5, 4, 3, 2, 1, following the INDIVIDUAL RATING RUBRIC and EVALUATION STEPS.

# Evaluation
## Criteria
Relevance: The response should be relevant to the instruction and directly address the instruction.

## Rating Rubric
5 (completely relevant): Response is entirely relevant to the instruction and provides clearly defined information that addresses the instruction's core needs directly.
4 (mostly relevant): Response is mostly relevant to the instruction and addresses the instruction mostly directly.
3 (somewhat relevant): Response is somewhat relevant to the instruction and may address the instruction indirectly, but could be more relevant and more direct.
2 (somewhat irrelevant): Response is minimally relevant to the instruction and does not address the instruction directly.
1 (irrelevant): Response is completely irrelevant to the instruction.

## Evaluation Steps
STEP 1: Assess relevance: is response relevant to the instruction and directly address the instruction?
STEP 2: Score based on the criteria and rubrics.

Give step by step explanations for your scoring, and only choose scores from 5, 4, 3, 2, 1.

# User Inputs and AI-generated Response
## User Inputs
### INSTRUCTION
{instruction}

### CONTEXT
{context}

## AI-generated Response
{response}
"""

relevance = PointwiseMetric(
    metric="relevance",
    metric_prompt_template=relevance_prompt_template,
)

helpfulness_prompt_template = """
# Instruction
You are a professional writing evaluator. Your job is to score writing responses according to pre-defined evaluation criteria.
You will be assessing question answering helpfulness, which measures the ability to provide important details when answering a question.
You will assign the writing response a score from 5, 4, 3, 2, 1, following the INDIVIDUAL RATING RUBRIC and EVALUATION STEPS.

# Evaluation
## Criteria
Helpfulness: The response is comprehensive with well-defined key details. The user would feel very satisfied with the content in a good response.

## Rating Rubric
5 (completely helpful): Response is useful and very comprehensive with well-defined key details to address the needs in the question and usually beyond what explicitly asked. The user would feel very satisfied with the content in the response.
4 (mostly helpful): Response is very relevant to the question, providing clearly defined information that addresses the question's core needs.  It may include additional insights that go slightly beyond the immediate question.  The user would feel quite satisfied with the content in the response.
3 (somewhat helpful): Response is relevant to the question and provides some useful content, but could be more relevant, well-defined, comprehensive, and/or detailed. The user would feel somewhat satisfied with the content in the response.
2 (somewhat unhelpful): Response is minimally relevant to the question and may provide some vaguely useful information, but it lacks clarity and detail. It might contain minor inaccuracies. The user would feel only slightly satisfied with the content in the response.
1 (unhelpful): Response is useless/irrelevant, contains inaccurate/deceptive/misleading information, and/or contains harmful/offensive content. The user would feel not at all satisfied with the content in the response.

## Evaluation Steps
STEP 1: Assess comprehensiveness: does the response provide specific, comprehensive, and clearly defined information for the user needs expressed in the question?
STEP 2: Assess relevance: When appropriate for the question, does the response exceed the question by providing relevant details and related information to contextualize content and help the user better understand the response.
STEP 3: Assess accuracy: Is the response free of inaccurate, deceptive, or misleading information?
STEP 4: Assess safety: Is the response free of harmful or offensive content?

Give step by step explanations for your scoring, and only choose scores from 5, 4, 3, 2, 1.

# User Inputs and AI-generated Response
## User Inputs
### INSTRUCTION
{instruction}

### CONTEXT
{context}

## AI-generated Response
{response}
"""

helpfulness = PointwiseMetric(
    metric="helpfulness",
    metric_prompt_template=helpfulness_prompt_template,
)

fulfillment_prompt_template = """
# Instruction
You are a professional writing evaluator. Your job is to score writing responses according to pre-defined evaluation criteria.
You will be assessing fulfillment, which measures the ability to follow instructions.
You will assign the writing response a score from 5, 4, 3, 2, 1, following the INDIVIDUAL RATING RUBRIC and EVALUATION STEPS.

# Evaluation
## Criteria
Instruction following: The response demonstrates a clear understanding of the instructions, satisfying all of the instruction's requirements.

## Rating Rubric
5 (complete fulfillment): Response addresses all aspects and adheres to all requirements of the instruction. The user would feel like their instruction was completely understood.
4 (good fulfillment): Response addresses most aspects and requirements of the instruction. It might miss very minor details or have slight deviations from requirements. The user would feel like their instruction was well understood.
3 (some fulfillment): Response does not address some minor aspects and/or ignores some requirements of the instruction. The user would feel like their instruction was partially understood.
2 (poor fulfillment): Response addresses some aspects of the instruction but misses key requirements or major components. The user would feel like their instruction was misunderstood in significant ways.
1 (no fulfillment): Response does not address the most important aspects of the instruction. The user would feel like their request was not at all understood.

## Evaluation Steps
STEP 1: Assess instruction understanding: Does the response address the intent of the instruction such that a user would not feel the instruction was ignored or misinterpreted by the response?
STEP 2: Assess requirements adherence: Does the response adhere to any requirements indicated in the instruction such as an explicitly specified word length, tone, format, or information that the response should include?

Give step by step explanations for your scoring, and only choose scores from 5, 4, 3, 2, 1.

# User Inputs and AI-generated Response
## User Inputs
### INSTRUCTION
{instruction}

## AI-generated Response
{response}
"""

fulfillment = PointwiseMetric(
    metric="fulfillment",
    metric_prompt_template=fulfillment_prompt_template,
)

eval_result = EvalTask(
    dataset=eval_dataset,
    metrics=[
        relevance,
        helpfulness,
        fulfillment,
    ],
    experiment=EXPERIMENT,
).evaluate()

notebook_utils.display_eval_result(eval_result)

new_eval_dataset = pd.DataFrame(
    {
        "prompt": [
            "Answer the question: " + question + " Context: " + context
            for question, context in zip(questions, retrieved_contexts)
        ],
        "response": generated_answers,
    }
)

# Run evaluation with new metric prompt template examples
eval_result = EvalTask(
    dataset=new_eval_dataset,
    metrics=[
        "question_answering_quality",
        "groundedness",
        "safety",
        "instruction_following",
    ],
    experiment=EXPERIMENT,
).evaluate()

notebook_utils.display_eval_result(eval_result)

eval_dataset = pd.DataFrame(
    {
        "prompt": [
            "Answer the question: " + question + " Context: " + context
            for question, context in zip(questions, retrieved_contexts)
        ],
    }
)

metric_name = "pairwise_text_quality"
pairwise_text_quality_result = EvalTask(
    dataset=eval_dataset,
    metrics=[
        PairwiseMetric(
            metric=metric_name,
            metric_prompt_template=MetricPromptTemplateExamples.get_prompt_template(
                metric_name
            ),
            # Specify baseline model for pairwise comparison
            baseline_model="gemini-2.5-flash",
        ),
    ],
    experiment=EXPERIMENT,
).evaluate(
    # Specify candidate model for pairwise comparison
    model="gemini-2.5-flash",
)

notebook_utils.display_eval_result(pairwise_text_quality_result)

notebook_utils.display_explanations(
    pairwise_text_quality_result, metrics=[metric_name], num=1
)

candidate_model_win_rate = round(
    pairwise_text_quality_result.summary_metrics[
        f"{metric_name}/candidate_model_win_rate"
    ]
    * 100
)
print(
    f"Win rate: Autorater prefers Candidate Model over Baseline Model {candidate_model_win_rate}% of time."
)

eval_dataset = pd.DataFrame(
    {
        "question": questions,
        "context": retrieved_contexts,
        "response": generated_answers,
        "baseline_model_response": baseline_answers,
    }
)

# Define an EvalTask with 2 example pairwise metrics
byor_pairwise_result = EvalTask(
    dataset=eval_dataset,
    metrics=[
        MetricPromptTemplateExamples.Pairwise.VERBOSITY,
        MetricPromptTemplateExamples.Pairwise.SAFETY,
    ],
    experiment=EXPERIMENT,
).evaluate(
    prompt_template="Answer the question: {question} Context: {context}",
    evaluation_service_qps=5,
)

notebook_utils.display_eval_result(byor_pairwise_result)

notebook_utils.display_explanations(
    byor_pairwise_result, metrics=["pairwise_verbosity"]
)

candidate_model_win_rate = round(
    byor_pairwise_result.summary_metrics["pairwise_verbosity/candidate_model_win_rate"]
    * 100
)
print(
    f"Win rate: Autorater prefers Candidate Model over Baseline Model {candidate_model_win_rate}% of time."
)