# 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 -q google-cloud-aiplatform google-cloud-bigquery bigframes pandas pandas-gbq google-genai

# You will see a notification of Colab crashing. It is the expected behavior.
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"
LOCATION = "us-central1"

# Generative model.
MODEL_ID = "gemini-2.5-flash"

# BigQuery tables.
BQ_DATASET_ID = "gemini_batch_predictions"
BQ_DATASET = f"{PROJECT_ID}.{BQ_DATASET_ID}"
FILES_TABLE = f"{BQ_DATASET_ID}.sports_files"
PROMPTS_TABLE = f"{BQ_DATASET}.temp_prompts"
TEXT_GENERATION_TABLE_PREFIX = f"{BQ_DATASET}.results"

# BigQuery views.
RESULTS_VIEW = f"{BQ_DATASET}.extraction_results"
EVALUATION_VIEW = f"{BQ_DATASET}.evaluation"

# File containing ground truth data in GCS.
BUCKET_NAME = "github-repo"
FOLDER = "generative-ai/gemini/evaluation/sports_files"
GCS_PREFIX = f"gs://{BUCKET_NAME}/{FOLDER}"
SPORTS_FILE = "sports_files.csv"

import datetime
import json
import time

import bigframes.pandas as bpd
import pandas as pd
import vertexai
from google import genai
from google.cloud import bigquery, storage
from google.genai import types
from vertexai.preview.batch_prediction import BatchPredictionJob

# BigQuery client.
bq_client = bigquery.Client(project=PROJECT_ID)

# Google Cloud Storage client.
storage_client = storage.Client()

# Initialize Vertex AI SDK.
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Set BigQuery Pandas options.
bpd.options.bigquery.project = PROJECT_ID
bpd.options.bigquery.location = LOCATION

def create_dataset(dataset_id: str, location: str) -> None:
    """Creates a BigQuery dataset in a location."""
    dataset = bigquery.Dataset(dataset_id)
    dataset.location = location

    dataset = bq_client.create_dataset(dataset, timeout=30)
    print(f"Created dataset {bq_client.project}.{dataset.dataset_id}")


def load_files_table_from_uri(files_table: str, uri: str) -> None:
    """Load ground truth into a BigQuery table from a GCS URI."""
    job_config = bigquery.LoadJobConfig(
        schema=[
            bigquery.SchemaField("path", "STRING"),
            bigquery.SchemaField("label", "STRING"),
        ],
        skip_leading_rows=1,
        source_format=bigquery.SourceFormat.CSV,
        write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE,
    )
    load_job = bq_client.load_table_from_uri(uri, files_table, job_config=job_config)
    load_job.result()

    destination_table = bq_client.get_table(files_table)
    print(f"Loaded {destination_table.num_rows} rows.")


create_dataset(dataset_id=BQ_DATASET, location=LOCATION)
load_files_table_from_uri(files_table=FILES_TABLE, uri=f"{GCS_PREFIX}/{SPORTS_FILE}")

ground_truth_df = bpd.read_gbq(FILES_TABLE)
ground_truth_df["path"][:2]

prompt = """\
- Classify the sport from the image below in one of the following categories:
* baseball
* basketball
* tennis
* volleyball

- Provide an answer in JSON format.

Example response:
{"sport": "baseball"}

- Image:
"""

def classify_image(model_id: str, prompt: str, gcs_prefix: str, blob_name: str) -> str:
    """Classifies an image."""
    client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

    response = client.models.generate_content(
        model=model_id,
        contents=[
            prompt,
            types.Part.from_uri(
                file_uri=f"{gcs_prefix}/{blob_name}", mime_type="image/jpeg"
            ),
        ],
        config=types.GenerateContentConfig(response_mime_type="application/json"),
    )
    return response.text


blob_name = ground_truth_df.iloc[0]["path"]
response = classify_image(
    model_id=MODEL_ID,
    prompt=prompt,
    gcs_prefix=GCS_PREFIX,
    blob_name=blob_name,
)
print(f"blob_name: {blob_name}")
print(f"response: {response}")

# Use current time as identifier of the evaluation.
now = datetime.datetime.now()
evaluation_ts = str(now)
evaluation_id = f"{now.year}_{now.month}_{now.day}_{now.hour}_{now.minute}"
json_file_name = f"/tmp/{evaluation_id}.json"

# Get URIs of the images from the ground truth table in BigQuery.
ground_truth_df = bpd.read_gbq(FILES_TABLE)

prompts_df = pd.DataFrame(
    [
        {
            "evaluation_ts": evaluation_ts,
            "evaluation_id": evaluation_id,
            "prompt_text": prompt,
            "gcs_uri": image_uri,
            "request": json.dumps(
                {
                    "contents": [
                        {
                            "role": "user",
                            "parts": [
                                {"text": prompt},
                                {
                                    "fileData": {
                                        "mimeType": "image/jpeg",
                                        "fileUri": f"{GCS_PREFIX}/{image_uri}",
                                    }
                                },
                            ],
                        }
                    ],
                    "generationConfig": {"responseMimeType": "application/json"},
                }
            ),
        }
        for image_uri in ground_truth_df["path"].values
    ]
)

# Save JSONL file
prompts_df.to_json(json_file_name, lines=True)

# Upload to BQ
prompts_df.to_gbq(PROMPTS_TABLE, PROJECT_ID)

table = bq_client.get_table(PROMPTS_TABLE)
print(
    f"Loaded {table.num_rows} rows and {len(table.schema)} columns to {PROMPTS_TABLE}"
)

# Define table to store results from Gemini Batch Prediction.
text_generation_table = f"{TEXT_GENERATION_TABLE_PREFIX}_{evaluation_id}"

# Create batch prediction job.
batch_job = BatchPredictionJob.submit(
    source_model=MODEL_ID,
    input_dataset=f"bq://{PROMPTS_TABLE}",
    output_uri_prefix=f"bq://{text_generation_table}",
)

# Refresh the job until complete
while not batch_job.has_ended:
    time.sleep(10)
    batch_job.refresh()

# Check if the job succeeds
if batch_job.has_succeeded:
    print("Job succeeded!")
else:
    print(f"Job failed: {batch_job.error}")

# Check the location of the output
print(f"Job output location: {batch_job.output_location}")

text_generation_df = bpd.read_gbq(text_generation_table)
for row in text_generation_df["response"][:5]:
    print(json.loads(row)["candidates"][0]["content"]["parts"][0]["text"])

def create_text_generation_view(
    text_generation_table_prefix: str, results_view: str
) -> None:
    """Creates a view of text extraction results."""
    view = bigquery.Table(results_view)

    view.view_query = rf"""
    SELECT
        evaluation_id,
        evaluation_ts,
        prompt_text,
        gcs_uri,
        JSON_VALUE(JSON_VALUE(response, '$.candidates[0].content.parts[0].text'), "$.sport") AS label
    FROM `{text_generation_table_prefix}_*`
    """

    # Make an API request to create the view.
    view = bq_client.create_table(view, exists_ok=False)
    print(f"Created {view.table_type}: {view.reference!s}")


create_text_generation_view(
    text_generation_table_prefix=TEXT_GENERATION_TABLE_PREFIX, results_view=RESULTS_VIEW
)

def create_evaluation_view(
    evaluation_view: str, files_table: str, results_view: str
) -> None:
    """Creates a view of experiment evaluation."""
    view = bigquery.Table(evaluation_view)

    view.view_query = f"""
      WITH t1 AS (
        SELECT
          e.evaluation_id,
          e.evaluation_ts,
          e.prompt_text,
          f.path,
          f.label,
          e.gcs_uri,
          f.label = e.label AS correct
        FROM `{files_table}` f
        JOIN `{results_view}` e
          ON f.path = e.gcs_uri
      )

      SELECT
        evaluation_id,
        evaluation_ts,
        prompt_text,
        path,
        label,
        correct
      FROM t1"""

    # Make an API request to create the view.
    view = bq_client.create_table(view, exists_ok=False)
    print(f"Created {view.table_type}: {view.reference!s}")


create_evaluation_view(
    evaluation_view=EVALUATION_VIEW, files_table=FILES_TABLE, results_view=RESULTS_VIEW
)