# 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 --quiet google-cloud-aiplatform

import sys

if "google.colab" in sys.modules:
    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"}


import vertexai

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

import json
import re
from typing import Any

import numpy as np
import pandas as pd
from vertexai.preview.evaluation import (
    AutoraterConfig,
    CustomOutputConfig,
    EvalTask,
    PointwiseMetric,
)

_TABLE_STYLE = [
    {
        "selector": "th",
        "props": [
            ("background-color", "#f2f2f2"),
            ("border", "1px solid gray"),
            ("color", "black"),
            ("font-size", "11pt"),
            ("text-align", "center"),
            ("word-break", "break-all"),
        ],
    },
    {
        "selector": "td",
        "props": [
            ("border", "1px solid gray"),
            ("color", "black"),
            ("min-width", "100px"),
            ("text-align", "center"),
        ],
    },
    {"selector": "tr:nth-child(even)", "props": [("background-color", "#f9f9f9")]},
    {"selector": "tr:nth-child(odd)", "props": [("background-color", "white")]},
    {"selector": "tr:hover", "props": [("background-color", "#94e6ff")]},
    {"selector": "td:hover", "props": [("background-color", "#ffffb3")]},
]


def parse_response_to_json(responses: list[str]) -> dict[str, Any]:
    response = re.sub(
        r"(.*```json|```.*)",
        "",
        responses[0].strip(),
    )
    sentences = []
    verdicts = []
    score = 0
    try:
        sentences = json.loads(response)
        for sentence in sentences:
            verdicts.append(sentence["label"])
        successful = 0
        for verdict in verdicts:
            if verdict == "supported":
                successful += 1
            elif verdict in ["unsupported", "contradictory"]:
                successful = 0
                break
        score = successful / len(verdicts)
    except Exception as e:
        print(f"Failed to parse JSON response: {e!s}")
    return {
        "sentence": sentences,
        "sentence_verdict": verdicts,
        "model_resp_score": score,
    }


def pretty_print_df(df: "pd.DataFrame", hide_columns: list[str]) -> "pd.Styler":
    styled_df = df.copy()
    for col in df.columns:
        if (
            isinstance(df[col][0], list)
            and df[col][0]
            and isinstance(df[col][0][0], dict)
        ):
            styled_df[col] = styled_df[col].apply(lambda x: _list_to_html_table(x))
    return (
        styled_df.style.hide(axis="index")
        .hide(subset=hide_columns, axis=1)
        .format({"groundedness/model_resp_score": "{:,.1f}"})
        .set_table_styles(_TABLE_STYLE)
    )


def _list_to_html_table(data: list[dict[str, Any]]) -> str:
    if not data:
        return "<i>No data to display.</i>"
    html_table = "<table style='border-collapse: collapse'><thead><tr>"
    # Extract headers from the first element
    for key in data[0].keys():
        html_table += f"<th>{key}</th>"
    html_table += "</tr></thead><tbody>"
    # Add rows
    for item in data:
        html_table += "<tr>"
        for value in item.values():
            html_table += f"<td>{value}</td>"
        html_table += "</tr>"
    html_table += "</tbody></table>"
    return html_table

GROUNDING_AUTORATER_PROMPT = """
You are a helpful and harmless AI assistant. You will be provided with a textual
context and a model-generated response.
Your task is to analyze the response sentence by sentence and classify each
sentence according to its relationship with the provided context.

**Instructions:**

1. **Decompose the response into individual sentences.**
2. **For each sentence, assign one of the following labels:**
    * **`supported`**: The sentence is entailed by the given context.  Provide a
      supporting excerpt from the context. The supporting excerpt must *fully*
      entail the sentence. If you need to cite multiple supporting excerpts,
      simply concatenate them.
    * **`unsupported`**: The sentence is not entailed by the given context. No
      excerpt is needed for this label.
    * **`contradictory`**: The sentence is falsified by the given context.
      Provide a contradicting excerpt from the context.
    * **`no_rad`**: The sentence does not require factual attribution (e.g.,
      opinions, greetings, questions, disclaimers).  No excerpt is needed for
      this label.
3. **For each label, provide a short rationale explaining your decision.**
   The rationale should be separate from the excerpt.
4. **Be very strict with your `supported` and `contradictory` decisions.**
   Unless you can find straightforward, indisputable evidence excerpts *in the
   context* that a sentence is `supported` or `contradictory`, consider it
   `unsupported`. You should not employ world knowledge unless it is truly
   trivial.

**Input Format:**

The input will consist of two parts, clearly separated:

* **Context:**  The textual context used to generate the response.
* **Response:** The model-generated response to be analyzed.

**Output Format:**

For each sentence in the response, output a JSON object with the following
fields:

* `"sentence"`: The sentence being analyzed.
* `"label"`: One of `supported`, `unsupported`, `contradictory`, or `no_rad`.
* `"rationale"`: A brief explanation for the assigned label.
* `"excerpt"`:  A relevant excerpt from the context. Only required for
  `supported` and `contradictory` labels.

Output each JSON object on a new line.

**Example:**

**Input:**

```
Context:
Apples are red fruits. Bananas are yellow fruits.

Response:
Apples are red. Bananas are green. Bananas are cheaper than apples. Enjoy your fruit!
```

**Output:**

[
{"sentence": "Apples are red.", "label": "supported", "rationale": "The context explicitly states that apples are red.", "excerpt": "Apples are red fruits."}
{"sentence": "Bananas are green.", "label": "contradictory", "rationale": "The context states that bananas are yellow, not green.", "excerpt": "Bananas are yellow fruits."}
{"sentence": "Bananas are cheaper than apples.", "label": "unsupported", "rationale": "The context does not mention the price of bananas or apples.", "excerpt": null}
{"sentence": "Enjoy your fruit!", "label": "no_rad", "rationale": "This is a general expression and does not require factual attribution.", "excerpt": null}
]

**Now, please analyze the following context and response:**

**User Query:**
{query}

**Context:**
{context}

**Response:**
{response}
"""

RESPONSE_PROMPT_TEMPLATE = """
Using only the information included in the context block, answer the user query
in 5 sentences or less.

**User Query:**
{query}

**Context:**
{context}
"""

grounded_metric = PointwiseMetric(
    metric="groundedness",
    metric_prompt_template=GROUNDING_AUTORATER_PROMPT,
    custom_output_config=CustomOutputConfig(
        return_raw_output=True,
        parsing_fn=parse_response_to_json,
    ),
)

# source: https://www.kaggle.com/datasets/deepmind/facts-grounding-examples/data?select=examples.csv
queries = [
    "What advantages do offline vs online retailers have?",
    "What are the requirements of OPM?",
    "How does virtual memory improve the efficiency of real physical memory (RAM) usage in computer systems?",
    "List the reasons that resulted in decreased emission of greenhouse gases from ethanol production.",
]
contexts = [
    """Due to savings in inventory costs, online retailing has a big advantage in selling
less popular items (the long tail) and in removing geographic barriers to purchase. Brynjolfsson
et al. (2003) estimated that the significantly increased assortment of books available online
increased consumer welfare by $700 million to a billion dollars in 2000. In comparing the online
sales of a clothing retailer with the catalog sales of the same item, Brynjolfsson et al. (2009)
showed that online sales of niche items were less sensitive to competition from offline stores
than from catalog sales because the online sales were skewed toward niche items. Brynjolfsson et
al. (2011) showed that online sales of niche items increased with recommendations and search tools,
indicating that these tools lowered search costs, making it easier for consumers to locate them.
In sum, because of the ability to handle more extensive inventories and provide search tools that
facilitate locating niche items, online retailing has a comparative advantage in selling less popular
items, translating into substantial benefits for consumers.
As noted above, online sellers have an advantage in facilitating a search for information on digital
attributes (including price). In contrast, offline sellers have an advantage in providing information
on non- digital attributes and providing faster delivery. This leads to the possibility that consumers
will search among both online and offline retailers.""",
    """As part of the assessment, S. 4043 would require OPM to explain whether each agency met its telework
goals and, if not, the actions being taken to identify and eliminate barriers to meeting them. The annual
report would also discuss additional steps that are planned by agencies to ensure telework oversight and
quality control and increase the utilization rates of office building space owned or leased by the agencies.
S. 4043 also requires the Office of Management and Budget (OMB), in consultation with GSA and the
Federal Real Property Council, to develop benchmarks and guidance for executive agencies to use when
calculating building utilization rates. S. 4043 would then require each executive agency head to establish
(1) a system to track office building space utilization rates consistent with that OMB guidance and (2)
indicators that measure the effects of telework policy on the management of real and personal property,
among other things.
S. 4043 would also require OPM to establish data standards to aid telework reporting requirements and
for automated telework tracking within payroll systems used by agencies. S. 4043 would require OPM, in
turn, to create an online tool that makes the standardized and reported data publicly available and would
allow OPM to use the online tool to fulfill its annual reporting requirements. For a more detailed
discussion of the bill's provisions on telework data standards, including office building utilization data,
see CRS Insight IN12352, Establishing Data Standards and Measuring Building Use: Select Provisions
of the Telework Transparency Act of 2024 (S. 4043).""",
    """Virtual memory is a computer system technique which gives an application program
the impression that it has contiguous working memory (an address space), while in fact it may be
physically fragmented and may even overflow on to disk storage. Systems that use this technique
make programming of large applications easier and use real physical memory (e.g.RAM) more
efficiently than those without virtual memory.
http://en.wikipedia.org/wiki/Virtual_memory
Page Fault: A page is a fixed-length block of memory that is used as a unit of transfer
between physical memory and external storage like a disk, and a page fault is an interrupt (or
exception) to the software raised by the hardware, when a program accesses a page that is
mapped in address space, but not loaded in physical memory.
http://en.wikipedia.org/wiki/Page_fault
Thrash is the term used to describe a degenerate situation on a computer where increasing
resources are used to do a decreasing amount of work. In this situation the system is
said to be thrashing. Usually it refers to two or more processes accessing a shared resource
repeatedly such that serious system performance degradation occurs because the system is
spending a disproportionate amount of time just accessing the shared resource. Resource
access time may generally be considered as wasted, since it does not contribute to the
advancement of any process. In modern computers, thrashing may occur in the paging system
(if there is not 'sufficient' physical memory or the disk access time is overly long), or in the
communications system (especially in conflicts over internal bus access), etc.
http://en.wikipedia.org/wiki/Thrash_(computer_science)""",
    """A new USDA report, titled 'A Life-Cycle Analysis of the Greenhouse Gas Emissions of Corn-Based
Ethanol,' finds that greenhouse gas emissions associated with producing corn-based ethanol in
the United States are about 43 percent lower than gasoline when measured on an energy equivalent
basis. Unlike other studies of greenhouse gas benefits, which relied on forecasts of future ethanol production
systems and expected impacts on the farm sector, this study reviewed how the industry and farm
sectors have performed over the past decade to assess the current greenhouse gas profile of corn-based ethanol.
The report shows that the reductions in greenhouse gas emissions were driven by a variety of improvements in
ethanol production, spanning from the corn field to the ethanol refinery. Farmers are producing corn
more efficiently and using conservation practices that reduce greenhouse gas emissions, including reduced tillage,
cover crops, and improved nitrogen management. Both corn yields and the efficiency of ethanol
production technologies are also improving.
Previous estimates of ethanol's greenhouse gas balance report lower efficiencies, largely due to anticipated
conversion of grasslands and forests to commodity production as a result of increased demand for corn
used in ethanol production. However, recent studies of international agricultural land use trends show
that since 2004, the primary land use change response of the world's farmers to rising commodity prices
has been to use available land resources more efficiently rather than to expand the amount of land used
for farming.""",
]

eval_dataset = pd.DataFrame(
    {
        "query": queries,
        "context": contexts,
    }
)

eval_task = EvalTask(
    dataset=eval_dataset,
    metrics=[grounded_metric],
    autorater_config=AutoraterConfig(sampling_count=1),
)

eval_result = eval_task.evaluate(
    model="gemini-2.5-flash",
    prompt_template=RESPONSE_PROMPT_TEMPLATE,
)
# Calculate overall score for metric.
np.mean(eval_result.metrics_table["groundedness/model_resp_score"])

pretty_print_df(
    eval_result.metrics_table, hide_columns=["prompt", "groundedness/sentence_verdict"]
)