from google.colab import auth

auth.authenticate_user()GOOGLE_CLOUD_PROJECT = "gcp-ml-172005"  # @param {type:"string"}
GOOGLE_CLOUD_REGION = "us-central1"

GCS_BUCKET_NAME = "cifar10-experimental-csp2"  # @param {type:"string"}
PIPELINE_NAME = "continuous-adaptation-for-data-changes-batch"  # @param {type:"string"}
PIPELINE_ROOT = "gs://{}/pipeline_root/{}".format(GCS_BUCKET_NAME, PIPELINE_NAME)
PIPELINE_LOCATION = f"{PIPELINE_ROOT}/{PIPELINE_NAME}_pipeline.json"

PUBSUB_TOPIC = f"trigger-{PIPELINE_NAME}"

SCHEDULER_JOB_NAME = f"scheduler-job-{PUBSUB_TOPIC}"

IMAGE_LOCATION_BUCKET = "batch-prediction-collection-3"  # @param {type:"string"}IMAGE_LOCATION_BUCKET_cloud_function_dep = "cloud_function/requirements.txt"%%writefile {_cloud_function_dep}

kfp==1.6.2
google-cloud-aiplatform
google-cloud-storage_cloud_function_file = "cloud_function/main.py"%%writefile {_cloud_function_file}

import os
import re
import json
import logging
import base64

from datetime import datetime

from kfp.v2.google.client import AIPlatformClient
from google.cloud import storage


def get_number_of_images(storage_client, bucket, latest_directory):
    blobs = storage_client.list_blobs(bucket, prefix=latest_directory)

    count = 0
    for blob in blobs:
        if blob.name.split(".")[-1] == "jpg":
            count = count + 1

    return count


def is_there_enough_images(storage_client, bucket, latest_directory, threshold):
    number_of_images = get_number_of_images(storage_client, bucket, latest_directory)
    print(f"number of images = {number_of_images}")
    return number_of_images >= threshold


def get_latest_directory(storage_client, bucket):
    blobs = storage_client.list_blobs(bucket)

    folders = list(
        set(
            [
                os.path.dirname(blob.name)
                for blob in blobs
                if bool(
                    re.match(
                        "[1-9][0-9][0-9][0-9]-[0-1][0-9]", os.path.dirname(blob.name)
                    )
                )
                is True
            ]
        )
    )

    folders.sort(key=lambda date: datetime.strptime(date, "%Y-%m"))
    print(folders[0])
    return folders[0]


def trigger_pipeline(event, context):
    # Parse the environment variables.
    project = os.getenv("PROJECT")
    region = os.getenv("REGION")
    gcs_pipeline_file_location = os.getenv("GCS_PIPELINE_FILE_LOCATION")
    gcs_image_bucket = os.getenv("GCS_IMAGE_BUCKET")

    print(project)
    print(region)
    print(gcs_pipeline_file_location)
    print(gcs_image_bucket)

    threshold = 100

    # Check if the pipeline file exists in the provided GCS Bucket.
    storage_client = storage.Client()
    latest_directory = get_latest_directory(storage_client, gcs_image_bucket)

    if is_there_enough_images(
        storage_client, gcs_image_bucket, latest_directory, threshold
    ):
        path_parts = gcs_pipeline_file_location.replace("gs://", "").split("/")
        pipeline_bucket = path_parts[0]
        pipeline_blob = "/".join(path_parts[1:])

        pipeline_bucket = storage_client.bucket(pipeline_bucket)
        blob = storage.Blob(bucket=pipeline_bucket, name=pipeline_blob)

        if not blob.exists(storage_client):
            raise ValueError(f"{gcs_pipeline_file_location} does not exist.")

        # Initialize Vertex AI API client and submit for pipeline execution.
        api_client = AIPlatformClient(project_id=project, region=region)

        response = api_client.create_run_from_job_spec(
            job_spec_path=gcs_pipeline_file_location,
            parameter_values={"data_gcs_prefix": latest_directory},
            enable_caching=True,
        )

        logging.info(response)ENV_VARS=f"""\
PROJECT={GOOGLE_CLOUD_PROJECT},\
REGION={GOOGLE_CLOUD_REGION},\
GCS_PIPELINE_FILE_LOCATION={PIPELINE_LOCATION},\
GCS_IMAGE_BUCKET={IMAGE_LOCATION_BUCKET}
"""

!echo {ENV_VARS}BUCKET = f'gs://{GCS_BUCKET_NAME}'
CLOUD_FUNCTION_NAME = f'trigger-{PIPELINE_NAME}-fn'

!gcloud functions deploy {CLOUD_FUNCTION_NAME} \
    --region={GOOGLE_CLOUD_REGION} \
    --trigger-topic={PUBSUB_TOPIC} \
    --runtime=python37 \
    --source=cloud_function\
    --entry-point=trigger_pipeline\
    --stage-bucket={BUCKET}\
    --update-env-vars={ENV_VARS}import IPython

cloud_fn_url = f"https://console.cloud.google.com/functions/details/{GOOGLE_CLOUD_REGION}/{CLOUD_FUNCTION_NAME}"
html = (
    f'See the Cloud Function details <a href="{cloud_fn_url}" target="_blank">here</a>.'
)
IPython.display.display(IPython.display.HTML(html))