# 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 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"}

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)
client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

import json

import pandas as pd
import vertexai
from vertexai import types

tool_call_valid_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    # Negative example: no tool calls expected
    if not target.get("tool_calls"):
        if not response.get("tool_calls"):
            return 1.0 # True negative
        return 0.0 # False positive

    # Tool call expected but not provided
    if not response.get("tool_calls"):
        return 0.0 # False negative

    target_call = target["tool_calls"][0]
    prediction_call = response["tool_calls"][0]

    if "name" in target_call and "name" in prediction_call:
        return 1.0
    return 0.0
"""

tool_call_valid_metric = types.Metric(
    name="tool_call_valid", remote_custom_function=tool_call_valid_code
)

tool_name_match_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    if not target.get("tool_calls"):
        return 1.0 if not response.get("tool_calls") else 0.0

    if not response.get("tool_calls"):
        return 0.0

    target_call = target["tool_calls"][0]
    prediction_call = response["tool_calls"][0]

    if target_call.get("name") == prediction_call.get("name"):
        return 1.0
    return 0.0
"""

tool_name_match_metric = types.Metric(
    name="tool_name_match", remote_custom_function=tool_name_match_code
)

tool_parameter_key_match_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    if not target.get("tool_calls"):
        return 1.0 if not response.get("tool_calls") else 0.0

    if not response.get("tool_calls"):
        return 0.0

    target_call = target["tool_calls"][0]
    prediction_call = response["tool_calls"][0]

    target_args = target_call.get("arguments", {})
    prediction_args = prediction_call.get("arguments", {})

    num_k_matches = sum(1 for k in target_args if k in prediction_args)
    num_unique_keys = len(set(target_args.keys()) | set(prediction_args.keys()))

    return float(num_k_matches) / float(num_unique_keys) if num_unique_keys > 0 else 1.0
"""

tool_parameter_kv_match_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    if not target.get("tool_calls"):
        return 1.0 if not response.get("tool_calls") else 0.0

    if not response.get("tool_calls"):
        return 0.0

    target_call = target["tool_calls"][0]
    prediction_call = response["tool_calls"][0]

    target_args = target_call.get("arguments", {})
    prediction_args = prediction_call.get("arguments", {})

    num_kv_matches = 0
    for k, v in target_args.items():
        if k in prediction_args and prediction_args[k] == v:
            num_kv_matches += 1

    num_unique_keys = len(set(target_args.keys()) | set(prediction_args.keys()))

    return float(num_kv_matches) / float(num_unique_keys) if num_unique_keys > 0 else 1.0
"""

tool_parameter_key_match_metric = types.Metric(
    name="tool_parameter_key_match",
    remote_custom_function=tool_parameter_key_match_code,
)

tool_parameter_kv_match_metric = types.Metric(
    name="tool_parameter_kv_match", remote_custom_function=tool_parameter_kv_match_code
)

# Create a sample evaluation dataset
dataset = pd.DataFrame(
    {
        "prompt": [
            "What is the weather in Tokyo?",
            "Tell me a joke.",
            "Book a flight to Paris tomorrow.",
        ],
        "reference": [
            json.dumps(
                {
                    "content": "",
                    "tool_calls": [
                        {"name": "get_weather", "arguments": {"location": "Tokyo"}}
                    ],
                }
            ),
            json.dumps({"content": "Here is a joke...", "tool_calls": []}),
            json.dumps(
                {
                    "content": "",
                    "tool_calls": [
                        {
                            "name": "book_flight",
                            "arguments": {"destination": "Paris", "date": "tomorrow"},
                        }
                    ],
                }
            ),
        ],
        "response": [
            json.dumps(
                {
                    "content": "",
                    "tool_calls": [
                        {"name": "get_weather", "arguments": {"location": "Tokyo"}}
                    ],
                }
            ),
            json.dumps({"content": "Why did the chicken...", "tool_calls": []}),
            json.dumps(
                {
                    "content": "",
                    "tool_calls": [
                        {"name": "book_flight", "arguments": {"destination": "Paris"}}
                    ],
                }
            ),
        ],
    }
)

eval_dataset = types.EvaluationDataset(eval_dataset_df=dataset)

# Evaluate using our custom metrics
eval_result = client.evals.evaluate(
    dataset=eval_dataset,
    metrics=[
        tool_call_valid_metric,
        tool_name_match_metric,
        tool_parameter_key_match_metric,
        tool_parameter_kv_match_metric,
    ],
)

eval_result.show()

# 1. Define Remote Custom Functions for Classification
tool_call_tp_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    ref_has_tool = len(target.get("tool_calls", [])) > 0
    pred_has_tool = len(response.get("tool_calls", [])) > 0

    if ref_has_tool and pred_has_tool:
        return 1.0
    return 0.0
"""

tool_call_fp_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    ref_has_tool = len(target.get("tool_calls", [])) > 0
    pred_has_tool = len(response.get("tool_calls", [])) > 0

    if not ref_has_tool and pred_has_tool:
        return 1.0
    return 0.0
"""

tool_call_fn_code = """
import json

def evaluate(instance: dict) -> float:
    ref_text = instance['reference']['contents']['gemini_contents'][0]['parts'][0]['text']
    pred_text = instance['response']['contents']['gemini_contents'][0]['parts'][0]['text']

    target = json.loads(ref_text)
    response = json.loads(pred_text)

    ref_has_tool = len(target.get("tool_calls", [])) > 0
    pred_has_tool = len(response.get("tool_calls", [])) > 0

    if ref_has_tool and not pred_has_tool:
        return 1.0
    return 0.0
"""

metric_tp = types.Metric(name="tool_call_tp", remote_custom_function=tool_call_tp_code)
metric_fp = types.Metric(name="tool_call_fp", remote_custom_function=tool_call_fp_code)
metric_fn = types.Metric(name="tool_call_fn", remote_custom_function=tool_call_fn_code)

# 2. Evaluate the dataset remotely
classification_eval_result = client.evals.evaluate(
    dataset=eval_dataset, metrics=[metric_tp, metric_fp, metric_fn]
)

# 3. Extract the Mean Scores
mean_tp = next(
    m.mean_score
    for m in classification_eval_result.summary_metrics
    if m.metric_name == "tool_call_tp"
)
mean_fp = next(
    m.mean_score
    for m in classification_eval_result.summary_metrics
    if m.metric_name == "tool_call_fp"
)
mean_fn = next(
    m.mean_score
    for m in classification_eval_result.summary_metrics
    if m.metric_name == "tool_call_fn"
)

# 4. Calculate overall Precision, Recall, and F1 mathematically
precision = mean_tp / (mean_tp + mean_fp) if (mean_tp + mean_fp) > 0 else 0.0
recall = mean_tp / (mean_tp + mean_fn) if (mean_tp + mean_fn) > 0 else 0.0
f1_score = (
    2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0
)

print("Overall Tool Usage Metrics (Calculated via Remote Custom Functions):")
print(f"Precision: {precision}")
print(f"Recall: {recall}")
print(f"F1 Score: {f1_score}")