# 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-genai jsonschema IPython==7.34.0 google-cloud-aiplatform 'pybind11>=2.12' deepdiff

import sys

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

    auth.authenticate_user()

# Create your own project and insert the project ID here ---->

import os

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

# Use the environment variable if the user doesn't provide Project ID.
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

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

EXPERIMENT_NAME = "eval-gemini-structured"

import vertexai
from google import genai

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)
vertexai.init(project=PROJECT_ID, location=LOCATION)

import json

import pandas as pd
from IPython.display import Image, Markdown, display
from deepdiff import DeepDiff
from google.genai.types import (
    GenerateContentConfig,
    Part,
)
from jsonschema import validate
from vertexai.evaluation import CustomMetric, EvalTask, notebook_utils

images = [
    {
        "image_url": "https://storage.googleapis.com/github-repo/generative-ai/gemini/evaluation/evaluate_gemini_structured_output/AcmeOrderForm.jpg",
        "image_uri": "gs://github-repo/generative-ai/gemini/evaluation/evaluate_gemini_structured_output/AcmeOrderForm.jpg",
        "image_type": "image/jpeg",
        "image_name": "Acme Order Form.jpg",
        "reference": {  # The Ground Truth
            "order_number": "98-X42-77A",
            "order_date": "2025-09-01",
            "customer_name": "WILE E. COYOTE (ESQ., PH.D, S.G.)",
            "customer_address": "HIGH MESA, CORNER OF X-MARK AND DETONATION CANYON, ANVIL FALLS, AZ",
            "line_items": [
                {
                    "item_description": "Jet Propelled Unicycle",
                    "quantity": 1,
                    "unit_price": 99.99,
                    "delivery_option": "Next Day",
                },
                {
                    "item_description": "Instant Hole Kit",
                    "quantity": 3,
                    "unit_price": 45.00,
                    "delivery_option": "Standard",
                },
                {
                    "item_description": "TNT High Explosives x24",
                    "quantity": 1,
                    "unit_price": 120.00,
                    "delivery_option": "Fast",
                },
                {
                    "item_description": "Super Magnet (XL)",
                    "quantity": 1,
                    "unit_price": 150.00,
                    "delivery_option": "Fast",
                },
                {
                    "item_description": "Rocket-Powered Roller skates",
                    "quantity": 2,
                    "unit_price": 79.99,
                    "delivery_option": "Next Day",
                },
            ],
        },
    },
    {
        "image_url": "https://storage.googleapis.com/github-repo/generative-ai/gemini/evaluation/evaluate_gemini_structured_output/EF0004.jpg",
        "image_uri": "gs://github-repo/generative-ai/gemini/evaluation/evaluate_gemini_structured_output/EF0004.jpg",
        "image_type": "image/jpeg",
        "image_name": "EF0004.jpg",
        "reference": {  # The Ground Truth
            "order_number": "EF0004",
            "order_date": "2025-10-26",
            "customer_name": "Elmer J. Fudd",
            "customer_address": "Happy Hunter's Hollow, Looney Tune Forest, CA",
            "line_items": [
                {
                    "item_description": "Silent Sneak Shoes",
                    "quantity": 1,
                    "unit_price": 35.99,
                    "delivery_option": "Standard",
                },
                {
                    "item_description": "Invisible Rabbit Trap",
                    "quantity": 2,
                    "unit_price": 75.00,
                    "delivery_option": "Standard",
                },
                {
                    "item_description": "Giant Butterfly Net",
                    "quantity": 1,
                    "unit_price": 49.50,
                    "delivery_option": "Fast",
                },
                {
                    "item_description": "Instant Camouflage Kit",
                    "quantity": 3,
                    "unit_price": 65.00,
                    "delivery_option": "Next Day",
                },
                {
                    "item_description": "Repellent Spray",
                    "quantity": 4,
                    "unit_price": 29.99,
                    "delivery_option": "Next Day",
                },
            ],
        },
    },
]

# Display the images using their public URLs

for image in images:
    print(image["image_name"])
    display(Image(url=image["image_url"], height=800))

# Define which models to compare

models = [
    # Gemini 2.0 family
    "gemini-2.0-flash",
    # Gemini 2.5 family
    "gemini-2.5-flash",
    "gemini-2.5-pro",
]

# Define the prompt and the structured output schema

prompt = """
    Analyze the attached scanned form and extract the information in the table in accordance with the schema.

    Provide the output in a clean JSON format.

    If any date field is formatted ambiguously, assume the dates are in dd/mm/yyyy format.

    If a field is blank, illegible, or cannot be found, return null for its value.

    If there are blank rows, do not include them in the output.

    If there is no image attached, return null for all fields.

"""

# Use structured output to ensure well formatted and consistent JSON output

schema = {
    "type": "object",
    "properties": {
        "order_number": {"type": "string"},
        "order_date": {
            "type": "string",
            "format": "date",  # Note: Enforces a full date output in the RFC 3339 format ("YYYY-MM-DD")
        },
        "customer_name": {"type": "string"},
        "customer_address": {"type": "string"},
        "line_items": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "item_description": {"type": "string"},
                    "quantity": {"type": "integer"},
                    "unit_price": {"type": "number"},
                    "delivery_option": {  # Note: We do not tell Gemini how to interpret the checkboxes as "Standard", "Fast" or "Next Day"
                        "type": "string"
                    },
                },
            },
        },
    },
}

generate_content_config = GenerateContentConfig(
    response_mime_type="application/json",
    response_schema=schema,
)

# Run the prompt for each model in `models` and each image in `images`, storing the output in `gemini_response`

gemini_response = {}
run_id = notebook_utils.generate_uuid(8)


for image_info in images:
    image = Part.from_uri(
        file_uri=image_info["image_uri"], mime_type=image_info["image_type"]
    )
    image_name = image_info["image_name"]

    gemini_response[image_name] = {}

    for model in models:
        run_name = f"{run_id}-{model}-{image_name}"

        response = client.models.generate_content(
            model=model, contents=[prompt, image], config=generate_content_config
        )

        response_json = json.dumps(response.parsed, indent=4)

        print("----------------------------------")
        print(f"{run_name}: \n{response_json}")
        gemini_response[image_name][model] = response_json

# Create the Evaluation Dataset

eval_dataset_rows = []
for image_info in images:
    image_name = image_info["image_name"]
    image_uri = image_info["image_uri"]
    image_type = image_info["image_type"]
    reference_str = json.dumps(
        image_info["reference"], indent=4
    )  # Convert the reference (ground truth) to pretty-printed JSON

    if image_name in gemini_response:
        models_data = gemini_response[image_name]
        for model_name, response_text in models_data.items():
            eval_dataset_rows.append(
                {
                    "model": model_name,
                    "prompt": prompt,  # The same prompt is used for all Gemini calls
                    "image": image_name,
                    "reference": reference_str,
                    "response": response_text,
                    "differences": DeepDiff(
                        json.loads(reference_str), json.loads(response_text)
                    ).pretty(),  # Uses the `deepdiff` library for identifying the differences between the response and the references
                }
            )

eval_dataset = pd.DataFrame(eval_dataset_rows)

# Define a custom evaluation metric to assess whether the response complies with the schema


def is_valid_schema(instance: dict[str, str]) -> dict[str, bool]:
    """Return 1 if the response complies with the schema, 0 if not."""
    response = instance["response"]

    try:
        validate(instance=json.loads(response), schema=schema)
    except Exception:
        return {"valid_schema": False}

    return {"valid_schema": True}


valid_schema = CustomMetric(name="valid_schema", metric_function=is_valid_schema)

# Define a `CustomMetric` for accuracy using the `deepdiff` library


def calculate_accuracy(instance: dict[str, str]) -> dict[str, float]:
    ref_json_string = instance["reference"]
    resp_json_string = instance["response"]

    try:
        reference_data = json.loads(ref_json_string)
        response_data = json.loads(resp_json_string)
    except json.JSONDecodeError:
        # If JSON is invalid or parsing fails, return 0 accuracy
        return {"accuracy": 0.0}

    # Use the deepdiff library to calculate the difference between the response and the reference (0 = exact match, 1 = )
    deep_distance = DeepDiff(
        reference_data, response_data, ignore_order=True, get_deep_distance=True
    ).get("deep_distance")

    accuracy = 1.0 - deep_distance

    return {"accuracy": accuracy}


accuracy = CustomMetric(name="accuracy", metric_function=calculate_accuracy)

# Define the evaluation task

extraction_eval_task = EvalTask(
    dataset=eval_dataset,
    metrics=[
        "exact_match",  # Exact match will only be 1 if the response is perfectly accurate, with no allowance for inconsistent JSON formatting. Hence, the custom `accuracy` metric is the better metric.
        valid_schema,
        accuracy,
    ],
    experiment=EXPERIMENT_NAME,
)

# Define the experiment & experiment run
run_id = notebook_utils.generate_uuid(8)

experiment_run_name = f"eval-{run_id}"

eval_result = extraction_eval_task.evaluate(experiment_run_name=experiment_run_name)

notebook_utils.display_eval_result(eval_result=eval_result)

# Let's get Gemini to analyse the results.

# Prepare the prompt for Gemini 2.5 Flash to summarize and analyze the results
summary_prompt = """
Analyze the following experiment results comparing Gemini models for extracted data from a scanned form.
The results include a summary table with overall metrics and row-based metrics, as well as the specific differences between the extracted data and the reference (ground truth).

Summarize the performance of each model based on the metrics provided (valid_schema, accuracy) from the summary table.
Analyze the detailed differences to understand the *types* of errors and mismatches occurring for each model.
Identify which models performed best and worst for each metric and based on the detailed error analysis.
Draw conclusions about the strengths and weaknesses of Gemini models for this specific tabular data extraction task, considering both the overall accuracy and the nature of the errors.
Consider the different versions of Gemini and how their performance varies.
Provide a clear and concise summary of the overall results, followed by key conclusions supported by observations from the detailed comparison.

Experiment Results Summary Table:
"""

# Convert the evaluation results summary and row-based metrics to a string format
# Assuming eval_result has a structure that can be converted to a readable string
try:
    # This will likely involve converting the DataFrames within eval_result to string
    eval_result_string = str(eval_result)
except Exception as e:
    eval_result_string = f"Could not convert evaluation results to string: {e}"
    print(eval_result_string)


# Concatenate the prompt and the summary table results
full_prompt = summary_prompt + eval_result_string

# Use Gemini 2.5 Flash to analyze the results
try:
    # Generate the response
    response = client.models.generate_content(
        model="gemini-2.5-flash", contents=full_prompt
    )

    # Display the summary and analysis from Gemini
    display(Markdown(response.text))

except Exception as e:
    print(f"An error occurred while calling Gemini: {e}")
    print(
        "Please ensure you have access to Gemini 2.5 Flash and your project/location settings are correct."
    )