import os

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# Google Cloud Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_GOOGLE_CLOUD_NOTEBOOK:
    USER_FLAG = "--user"

# Install Python package dependencies.
print("Installing libraries")
! pip3 install {USER_FLAG} --quiet google-cloud-pipeline-components kfp
! pip3 install {USER_FLAG} --quiet --upgrade google-cloud-aiplatform google-cloud-bigquery

"""
Explanation: Experimenting BQML and AutoML with Vertex AI
Overview
This notebook demonstrates how to use Vertex AI Pipelines to rapid prototype a model using both AutoML and BQML, do an evaluation comparison, for a baseline, before progressing to a custom model.
Learning objectives
In this notebook, you learn how to use Vertex AI Predictions for rapid prototyping a model.
This notebook uses the following Google Cloud ML services:

Vertex AI Pipelines
Vertex AI AutoML
Vertex AI BigQuery ML
Google Cloud Pipeline Components

The steps performed include:

Create a BigQuery and Vertex AI training dataset.
Train a BigQuery ML and AutoML model.
Extract evaluation metrics from the BigQueryML and AutoML models.
Select and deploy the best trained model.
Test the deployed model infrastructure.
Run the pipeline.

Each learning objective will correspond to a #TODO in this student lab notebook -- try to complete this notebook first and then review the solution notebook.
<img src="https://storage.googleapis.com/rafacarv-public-bucket-do-not-delete/abalone/automl_and_bqml.png" />
Dataset
The Abalone Dataset
<img src="https://storage.googleapis.com/rafacarv-public-bucket-do-not-delete/abalone/dataset.png" />
<p>Dataset Credits</p>
<p>Dua, D. and Graff, C. (2019). UCI Machine Learning Repository <a href="http://archive.ics.uci.edu/ml">http://archive.ics.uci.edu/ml</a>. Irvine, CA: University of California, School of Information and Computer Science.</p>

<p><a href="https://archive.ics.uci.edu/ml/datasets/abalone">Direct link</a></p>

Attribute Information:
<p>Given is the attribute name, attribute type, the measurement unit, and a brief description. The number of rings is the value to predict: either as a continuous value or as a classification problem.</p>

<body>
    <table>
        <tr>
            <th>Name</th>
            <th>Data Type</th>
            <th>Measurement Unit</th>
            <th>Description</th>
        </tr>
        <tr>
            <td>Sex</td>
            <td>nominal</td>
            <td>--</td>
            <td>M, F, and I (infant)</td>
        </tr>
        <tr>
            <td>Length</td>
            <td>continuous</td>
            <td>mm</td>
            <td>Longest shell measurement</td>
        </tr>
        <tr>
            <td>Diameter</td>
            <td>continuous</td>
            <td>mm</td>
            <td>perpendicular to length</td>
        </tr>
        <tr>
            <td>Height</td>
            <td>continuous</td>
            <td>mm</td>
            <td>with meat in shell</td>
        </tr>
        <tr>
            <td>Whole weight</td>
            <td>continuous</td>
            <td>grams</td>
            <td>whole abalone</td>
        </tr>
        <tr>
            <td>Shucked weight</td>
            <td>continuous</td>
            <td>grams</td>
            <td>weight of meat</td>
        </tr>
        <tr>
            <td>Viscera weight</td>
            <td>continuous</td>
            <td>grams</td>
            <td>gut weight (after bleeding)</td>
        </tr>
        <tr>
            <td>Shell weight</td>
            <td>continuous</td>
            <td>grams</td>
            <td>after being dried</td>
        </tr>
        <tr>
            <td>Rings</td>
            <td>integer</td>
            <td>--</td>
            <td>+1.5 gives the age in years</td>
        </tr>
    </table>
</body>

Install additional packages
End of explanation
"""


# Automatically restart kernel after installs
import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

"""
Explanation: Please ignore any incompatibility warnings and errors.
End of explanation
"""


import os

PROJECT_ID = ""

# Get your Google Cloud project ID from gcloud
if not os.getenv("IS_TESTING"):
    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID: ", PROJECT_ID)

"""
Explanation: Set your project ID
If you don't know your project ID, you may be able to get your project ID using gcloud.
End of explanation
"""


if PROJECT_ID == "" or PROJECT_ID is None:
    PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

! gcloud config set project $PROJECT_ID

"""
Explanation: Otherwise, set your project ID here.
End of explanation
"""


from datetime import datetime

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")

"""
Explanation: Timestamp
If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial.
End of explanation
"""


BUCKET_NAME = "[your-bucket-name]"  # @param {type:"string"}
REGION = "[your-region]"  # @param {type:"string"}

if BUCKET_NAME == "" or BUCKET_NAME is None or BUCKET_NAME == "[your-bucket-name]":
    BUCKET_NAME = PROJECT_ID + "aip-" + TIMESTAMP
BUCKET_URI = f"gs://{BUCKET_NAME}"

if REGION == "[your-region]":
    REGION = "us-central1"

"""
Explanation: Create a Cloud Storage bucket
The following steps are required, regardless of your notebook environment.
When you submit a training job using the Cloud SDK, you upload a Python package
containing your training code to a Cloud Storage bucket. Vertex AI runs
the code from this package. In this notebook, Vertex AI also saves the
trained model that results from your job in the same bucket. Using this model artifact, you can then
create Vertex AI model and endpoint resources in order to serve
online predictions.
Set the name of your Cloud Storage bucket below. It must be unique across all
Cloud Storage buckets.
You may also change the REGION variable, which is used for operations
throughout the rest of this notebook. It is suggested that you choose a region where Vertex AI services are
available.
End of explanation
"""


! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI

"""
Explanation: Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.
End of explanation
"""


! gsutil ls -al $BUCKET_URI

"""
Explanation: Finally, validate access to your Cloud Storage bucket by examining its contents:
End of explanation
"""


SERVICE_ACCOUNT = "[your-service-account]"  # @param {type:"string"}

if (
    SERVICE_ACCOUNT == ""
    or SERVICE_ACCOUNT is None
    or SERVICE_ACCOUNT == "[your-service-account]"
):
    # Get your GCP project id from gcloud
    shell_output = !gcloud auth list 2>/dev/null
    SERVICE_ACCOUNT = shell_output[2].replace("*", "").strip()
    print("Service Account:", SERVICE_ACCOUNT)

"""
Explanation: Service Account
If you don't know your service account, try to get your service account using gcloud command by executing the second cell below.
End of explanation
"""


! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI

! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI

"""
Explanation: Set service account access for Vertex AI Pipelines
Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account.
End of explanation
"""


import sys
from typing import NamedTuple

from google.cloud import aiplatform as vertex
from google.cloud import bigquery
from google_cloud_pipeline_components import \
    aiplatform as vertex_pipeline_components
from google_cloud_pipeline_components.experimental import \
    bigquery as bq_components
from kfp import dsl
from kfp.v2 import compiler
from kfp.v2.dsl import Artifact, Input, Metrics, Output, component

"""
Explanation: Required imports
End of explanation
"""


PIPELINE_JSON_PKG_PATH = "rapid_prototyping.json"
PIPELINE_ROOT = f"gs://{BUCKET_NAME}/pipeline_root"
DATA_FOLDER = f"{BUCKET_NAME}/data"

RAW_INPUT_DATA = f"gs://{DATA_FOLDER}/abalone.csv"
BQ_DATASET = "j90wipxexhrgq3cquanc5"  # @param {type:"string"}
BQ_LOCATION = "US"  # @param {type:"string"}
BQ_LOCATION = BQ_LOCATION.upper()
BQML_EXPORT_LOCATION = f"gs://{BUCKET_NAME}/artifacts/bqml"

DISPLAY_NAME = "rapid-prototyping"
ENDPOINT_DISPLAY_NAME = f"{DISPLAY_NAME}_endpoint"

image_prefix = REGION.split("-")[0]
BQML_SERVING_CONTAINER_IMAGE_URI = (
    f"{image_prefix}-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest"
)

if os.getenv("IS_TESTING"):
    !gcloud --quiet components install beta
    !gcloud --quiet components update
!gcloud config set project $PROJECT_ID
!gcloud config set ai/region $REGION

"""
Explanation: Determine some project and pipeline variables
Instructions:
- Make sure the GCS bucket and the BigQuery Dataset do not exist. This script may delete any existing content.
- Your bucket must be on the same region as your Vertex AI resources.
- BQ region can be US or EU;
- Make sure your preferred Vertex AI region is supported [link].
End of explanation
"""


! gsutil cp gs://cloud-samples-data/vertex-ai/community-content/datasets/abalone/abalone.data {RAW_INPUT_DATA}

"""
Explanation: Downloading the data
The cell below will download the dataset into a CSV file and save it in GCS
End of explanation
"""


@component(base_image="python:3.9", packages_to_install=["google-cloud-bigquery"])
def import_data_to_bigquery(
    project: str,
    bq_location: str,
    bq_dataset: str,
    gcs_data_uri: str,
    raw_dataset: Output[Artifact],
    table_name_prefix: str = "abalone",
):
    from google.cloud import bigquery

    # Construct a BigQuery client object.
    client = # TODO 1: Your code goes here

    def load_dataset(gcs_uri, table_id):
        job_config = bigquery.LoadJobConfig(
            schema=[
                bigquery.SchemaField("Sex", "STRING"),
                bigquery.SchemaField("Length", "NUMERIC"),
                bigquery.SchemaField("Diameter", "NUMERIC"),
                bigquery.SchemaField("Height", "NUMERIC"),
                bigquery.SchemaField("Whole_weight", "NUMERIC"),
                bigquery.SchemaField("Shucked_weight", "NUMERIC"),
                bigquery.SchemaField("Viscera_weight", "NUMERIC"),
                bigquery.SchemaField("Shell_weight", "NUMERIC"),
                bigquery.SchemaField("Rings", "NUMERIC"),
            ],
            skip_leading_rows=1,
            # The source format defaults to CSV, so the line below is optional.
            source_format=bigquery.SourceFormat.CSV,
        )
        print(f"Loading {gcs_uri} into {table_id}")
        load_job = client.load_table_from_uri(
            gcs_uri, table_id, job_config=job_config
        )  # Make an API request.

        load_job.result()  # Waits for the job to complete.
        destination_table = client.get_table(table_id)  # Make an API request.
        print("Loaded {} rows.".format(destination_table.num_rows))

    def create_dataset_if_not_exist(bq_dataset_id, bq_location):
        print(
            "Checking for existence of bq dataset. If it does not exist, it creates one"
        )
        dataset = bigquery.Dataset(bq_dataset_id)
        dataset.location = bq_location
        dataset = client.create_dataset(dataset, exists_ok=True, timeout=300)
        print(f"Created dataset {dataset.full_dataset_id} @ {dataset.location}")

    bq_dataset_id = f"{project}.{bq_dataset}"
    create_dataset_if_not_exist(bq_dataset_id, bq_location)

    raw_table_name = f"{table_name_prefix}_raw"
    table_id = f"{project}.{bq_dataset}.{raw_table_name}"
    print("Deleting any tables that might have the same name on the dataset")
    client.delete_table(table_id, not_found_ok=True)
    print("will load data to table")
    load_dataset(gcs_data_uri, table_id)

    raw_dataset_uri = f"bq://{table_id}"
    raw_dataset.uri = raw_dataset_uri

"""
Explanation: Pipeline Components
Import to BQ
This component takes the csv file and imports it to a table in BigQuery. If the dataset does not exist, it will be created. If a table with the same name already exists, it will be deleted and recreated
End of explanation
"""


@component(
    base_image="python:3.9",
    packages_to_install=["google-cloud-bigquery"],
)  # pandas, pyarrow and fsspec required to export bq data to csv
def split_datasets(
    raw_dataset: Input[Artifact],
    bq_location: str,
) -> NamedTuple(
    "bqml_split",
    [
        ("dataset_uri", str),
        ("dataset_bq_uri", str),
        ("test_dataset_uri", str),
    ],
):

    from collections import namedtuple

    from google.cloud import bigquery

    raw_dataset_uri = raw_dataset.uri
    table_name = raw_dataset_uri.split("bq://")[-1]
    print(table_name)
    raw_dataset_uri = table_name.split(".")
    print(raw_dataset_uri)
    project = raw_dataset_uri[0]
    bq_dataset = raw_dataset_uri[1]
    bq_raw_table = raw_dataset_uri[2]

    client = bigquery.Client(project=project, location=bq_location)

    def split_dataset(table_name_dataset):
        training_dataset_table_name = f"{project}.{bq_dataset}.{table_name_dataset}"
        split_query = f"""
        CREATE OR REPLACE TABLE
            `{training_dataset_table_name}`
           AS
        SELECT
          Sex,
          Length,
          Diameter,
          Height,
          Whole_weight,
          Shucked_weight,
          Viscera_weight,
          Shell_weight,
          Rings,
            CASE(ABS(MOD(FARM_FINGERPRINT(TO_JSON_STRING(f)), 10)))
              WHEN 9 THEN 'TEST'
              WHEN 8 THEN 'VALIDATE'
              ELSE 'TRAIN' END AS split_col
        FROM
          `{project}.{bq_dataset}.abalone_raw` f
        """
        dataset_uri = f"{project}.{bq_dataset}.{bq_raw_table}"
        print("Splitting the dataset")
        query_job = client.query(split_query)  # Make an API request.
        query_job.result()
        print(dataset_uri)
        print(split_query.replace("\n", " "))
        return training_dataset_table_name

    def create_test_view(training_dataset_table_name, test_view_name="dataset_test"):
        view_uri = f"{project}.{bq_dataset}.{test_view_name}"
        query = f"""
             CREATE OR REPLACE VIEW `{view_uri}` AS SELECT
          Sex,
          Length,
          Diameter,
          Height,
          Whole_weight,
          Shucked_weight,
          Viscera_weight,
          Shell_weight,
          Rings 
          FROM `{training_dataset_table_name}`  f
          WHERE 
          f.split_col = 'TEST'
          """
        print(f"Creating view for --> {test_view_name}")
        print(query.replace("\n", " "))
        query_job = client.query(query)  # Make an API request.
        query_job.result()
        return view_uri

    table_name_dataset = "dataset"

    dataset_uri = split_dataset(table_name_dataset)
    test_dataset_uri = create_test_view(dataset_uri)
    dataset_bq_uri = "bq://" + dataset_uri

    print(f"dataset: {dataset_uri}")

    result_tuple = namedtuple(
        "bqml_split",
        ["dataset_uri", "dataset_bq_uri", "test_dataset_uri"],
    )
    return result_tuple(
        dataset_uri=str(dataset_uri),
        dataset_bq_uri=str(dataset_bq_uri),
        test_dataset_uri=str(test_dataset_uri),
    )

"""
Explanation: Split Datasets
Splits the dataset in 3 slices:
- TRAIN
- EVALUATE
- TEST
AutoML and BigQuery ML use different nomenclatures for data splits:
BQML
How BQML splits the data: link
AutoML
How AutoML splits the data: link
<ul>
    <li>Model trials
<p>The training set is used to train models with different preprocessing, architecture, and hyperparameter option combinations. These models are evaluated on the validation set for quality, which guides the exploration of additional option combinations. The best parameters and architectures determined in the parallel tuning phase are used to train two ensemble models as described below.</p></li>

<li>Model evaluation
<p>
Vertex AI trains an evaluation model, using the training and validation sets as training data. Vertex AI generates the final model evaluation metrics on this model, using the test set. This is the first time in the process that the test set is used. This approach ensures that the final evaluation metrics are an unbiased reflection of how well the final trained model will perform in production.</p></li>

<li>Serving model
<p>A model is trained with the training, validation, and test sets, to maximize the amount of training data. This model is the one that you use to request predictions.</p></li>
End of explanation
"""


def _query_create_model(
    project_id: str,
    bq_dataset: str,
    training_data_uri: str,
    model_name: str = "linear_regression_model_prototyping",
):
    model_uri = f"{project_id}.{bq_dataset}.{model_name}"

    model_options = """OPTIONS
      ( MODEL_TYPE='LINEAR_REG',
        input_label_cols=['Rings'],
         DATA_SPLIT_METHOD='CUSTOM',
        DATA_SPLIT_COL='split_col'
        )
        """
    query = f"""
    CREATE OR REPLACE MODEL
      `{model_uri}`
      {model_options}
     AS
    SELECT
      Sex,
      Length,
      Diameter,
      Height,
      Whole_weight,
      Shucked_weight,
      Viscera_weight,
      Shell_weight,
      Rings,
      CASE(split_col)
        WHEN 'TEST' THEN TRUE
      ELSE
      FALSE
    END
      AS split_col
    FROM
      `{training_data_uri}`;
    """

    print(query.replace("\n", " "))

    return query

"""
Explanation: Train BQML Model
For this demo, you will use a simple linear regression model on BQML. However, you can be creative with other model architectures, such as Deep Neural Networks, XGboost, Logistic Regression, etc.
For a full list of models supported by BQML, look here: End-to-end user journey for each model.
As pointed out before, BQML and AutoML use different split terminologies, so you do an adaptation of the <i>split_col</i> column directly on the SELECT portion of the CREATE model query:

When the value of DATA_SPLIT_METHOD is 'CUSTOM', the corresponding column should be of type BOOL. The rows with TRUE or NULL values are used as evaluation data. Rows with FALSE values are used as training data.
End of explanation
"""


@component(base_image="python:3.9")
def interpret_bqml_evaluation_metrics(
    bqml_evaluation_metrics: Input[Artifact], metrics: Output[Metrics]
) -> dict:
    import math

    metadata = bqml_evaluation_metrics.metadata
    for r in metadata["rows"]:

        rows = r["f"]
        schema = metadata["schema"]["fields"]

        output = {}
        for metric, value in zip(schema, rows):
            metric_name = metric["name"]
            val = float(value["v"])
            output[metric_name] = val
            metrics.log_metric(metric_name, val)
            if metric_name == "mean_squared_error":
                rmse = math.sqrt(val)
                metrics.log_metric("root_mean_squared_error", rmse)

    metrics.log_metric("framework", "BQML")

    print(output)

"""
Explanation: Interpret BQML Model Evaluation
When you do Hyperparameter tuning on the model creation query, the output of the pre-built component BigqueryEvaluateModelJobOp will be a table with the metrics obtained by BQML when training the model. In your BigQuery console, they look like the image below. You need to access them programmatically so you can compare them to the AutoML model. 
<img src="https://storage.googleapis.com/rafacarv-public-bucket-do-not-delete/abalone/bqml-evaluate.png?">
The cell below shows you an example of how this can be done. BQML does not give you a root mean squared error to the list of metrics, so we're manually adding it to the metrics dictionary. For more information about the output, please check BQML's documentation.
End of explanation
"""


# Inspired by Andrew Ferlitsch's work on https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_automl_pipeline_components.ipynb


@component(
    base_image="python:3.9",
    packages_to_install=[
        "google-cloud-aiplatform",
    ],
)
def interpret_automl_evaluation_metrics(
    region: str, model: Input[Artifact], metrics: Output[Metrics]
):
    """'
    For a list of available regression metrics, go here: gs://google-cloud-aiplatform/schema/modelevaluation/regression_metrics_1.0.0.yaml.

    More information on available metrics for different types of models: https://cloud.google.com/vertex-ai/docs/predictions/online-predictions-automl
    """

    import google.cloud.aiplatform.gapic as gapic

    # Get a reference to the Model Service client
    client_options = {"api_endpoint": f"{region}-aiplatform.googleapis.com"}

    model_service_client = gapic.ModelServiceClient(client_options=client_options)

    model_resource_name = model.metadata["resourceName"]

    model_evaluations = model_service_client.list_model_evaluations(
        parent=model_resource_name
    )
    # List the model evaluations.
    model_evaluation = # TODO 2: Your code goes here

    available_metrics = [
        "meanAbsoluteError",
        "meanAbsolutePercentageError",
        "rSquared",
        "rootMeanSquaredError",
        "rootMeanSquaredLogError",
    ]
    output = dict()
    for x in available_metrics:
        val = model_evaluation.metrics.get(x)
        output[x] = val
        metrics.log_metric(str(x), float(val))

    metrics.log_metric("framework", "AutoML")
    print(output)

"""
Explanation: Interpret AutoML Model Evaluation
Similar to BQML, AutoML also generates metrics during its model creation. These can be accessed in the UI, as seen below:
<img src="https://storage.googleapis.com/rafacarv-public-bucket-do-not-delete/abalone/automl-evaluate.png" />
Since you don't have a pre-built-component to access these metrics programmatically, you can use the Vertex AI GAPIC (Google API Compiler), which auto-generates low-level gRPC interfaces to the service.
End of explanation
"""


@component(base_image="python:3.9")
def select_best_model(
    metrics_bqml: Input[Metrics],
    metrics_automl: Input[Metrics],
    thresholds_dict_str: str,
    best_metrics: Output[Metrics],
    reference_metric_name: str = "rmse",
) -> NamedTuple(
    "Outputs",
    [
        ("deploy_decision", str),
        ("best_model", str),
        ("metric", float),
        ("metric_name", str),
    ],
):
    import json
    from collections import namedtuple

    best_metric = float("inf")
    best_model = None

    # BQML and AutoML use different metric names.
    metric_possible_names = []

    if reference_metric_name == "mae":
        metric_possible_names = ["meanAbsoluteError", "mean_absolute_error"]
    elif reference_metric_name == "rmse":
        metric_possible_names = ["rootMeanSquaredError", "root_mean_squared_error"]

    metric_bqml = float("inf")
    metric_automl = float("inf")
    print(metrics_bqml.metadata)
    print(metrics_automl.metadata)
    for x in metric_possible_names:

        try:
            metric_bqml = metrics_bqml.metadata[x]
            print(f"Metric bqml: {metric_bqml}")
        except:
            print(f"{x} does not exist int the BQML dictionary")

        try:
            metric_automl = metrics_automl.metadata[x]
            print(f"Metric automl: {metric_automl}")
        except:
            print(f"{x} does not exist on the AutoML dictionary")

    # Change the condition if higher is better.
    print(f"Comparing BQML ({metric_bqml}) vs AutoML ({metric_automl})")
    if metric_bqml <= metric_automl:
        best_model = "bqml"
        best_metric = metric_bqml
        best_metrics.metadata = metrics_bqml.metadata
    else:
        best_model = "automl"
        best_metric = metric_automl
        best_metrics.metadata = metrics_automl.metadata

    thresholds_dict = json.loads(thresholds_dict_str)
    deploy = False

    # Change the condition if higher is better.
    # TODO 3: Your code goes here

    if deploy:
        deploy_decision = "true"
    else:
        deploy_decision = "false"

    print(f"Which model is best? {best_model}")
    print(f"What metric is being used? {reference_metric_name}")
    print(f"What is the best metric? {best_metric}")
    print(f"What is the threshold to deploy? {thresholds_dict_str}")
    print(f"Deploy decision: {deploy_decision}")

    Outputs = namedtuple(
        "Outputs", ["deploy_decision", "best_model", "metric", "metric_name"]
    )

    return Outputs(
        deploy_decision=deploy_decision,
        best_model=best_model,
        metric=best_metric,
        metric_name=reference_metric_name,
    )

"""
Explanation: Model Selection
Now that you have evaluated the models independently, you are going to move forward with only one of them. This election will be done based on the model evaluation metrics gathered in the previous steps.
Bear in mind that BQML and AutoML use different evaluation metric names, hence you had to do a mapping of these different nomenclatures.
End of explanation
"""


@component(base_image="python:3.9", packages_to_install=["google-cloud-aiplatform"])
def validate_infrastructure(
    endpoint: Input[Artifact],
) -> NamedTuple(
    "validate_infrastructure_output", [("instance", str), ("prediction", float)]
):
    import json
    from collections import namedtuple

    from google.cloud import aiplatform
    from google.protobuf import json_format
    from google.protobuf.struct_pb2 import Value

    def treat_uri(uri):
        return uri[uri.find("projects/") :]

    def request_prediction(endp, instance):
        instance = json_format.ParseDict(instance, Value())
        instances = [instance]
        parameters_dict = {}
        parameters = json_format.ParseDict(parameters_dict, Value())
        response = endp.predict(instances=instances, parameters=parameters)
        print("deployed_model_id:", response.deployed_model_id)
        print("predictions: ", response.predictions)
        # The predictions are a google.protobuf.Value representation of the model's predictions.
        predictions = response.predictions

        for pred in predictions:
            if type(pred) is dict and "value" in pred.keys():
                # AutoML predictions
                prediction = pred["value"]
            elif type(pred) is list:
                # BQML Predictions return different format
                prediction = pred[0]
            return prediction

    endpoint_uri = endpoint.uri
    treated_uri = treat_uri(endpoint_uri)

    instance = {
        "Sex": "M",
        "Length": 0.33,
        "Diameter": 0.255,
        "Height": 0.08,
        "Whole_weight": 0.205,
        "Shucked_weight": 0.0895,
        "Viscera_weight": 0.0395,
        "Shell_weight": 0.055,
    }
    instance_json = json.dumps(instance)
    print("Will use the following instance: " + instance_json)

    endpoint = aiplatform.Endpoint(treated_uri)

    # Make a simple prediction
    prediction = # TODO 4: Your code goes here
    result_tuple = namedtuple(
        "validate_infrastructure_output", ["instance", "prediction"]
    )

    return result_tuple(instance=str(instance_json), prediction=float(prediction))

"""
Explanation: Validate Infrastructure
Once the best model has been deployed, you will validate the endpoint by making a simple prediction to it.
End of explanation
"""


pipeline_params = {
    "project": PROJECT_ID,
    "region": REGION,
    "gcs_input_file_uri": RAW_INPUT_DATA,
    "bq_dataset": BQ_DATASET,
    "bq_location": BQ_LOCATION,
    "bqml_model_export_location": BQML_EXPORT_LOCATION,
    "bqml_serving_container_image_uri": BQML_SERVING_CONTAINER_IMAGE_URI,
    "endpoint_display_name": ENDPOINT_DISPLAY_NAME,
    "thresholds_dict_str": '{"rmse": 2.5}',
}

@dsl.pipeline(name=DISPLAY_NAME, description="Rapid Prototyping")
def train_pipeline(
    project: str,
    gcs_input_file_uri: str,
    region: str,
    bq_dataset: str,
    bq_location: str,
    bqml_model_export_location: str,
    bqml_serving_container_image_uri: str,
    endpoint_display_name: str,
    thresholds_dict_str: str,
):
    # Imports data to BigQuery using a custom component.
    import_data_to_bigquery_op = import_data_to_bigquery(
        project, bq_location, bq_dataset, gcs_input_file_uri
    )
    raw_dataset = import_data_to_bigquery_op.outputs["raw_dataset"]

    # Splits the BQ dataset using a custom component.
    split_datasets_op = split_datasets(raw_dataset, bq_location=bq_location)

    # Generates the query to create a BQML using a static function.
    create_model_query = _query_create_model(
        project, bq_dataset, split_datasets_op.outputs["dataset_uri"]
    )

    # Builds BQML model using pre-built-component.
    bqml_create_op = bq_components.BigqueryCreateModelJobOp(
        project=project, location=bq_location, query=create_model_query
    )
    bqml_model = bqml_create_op.outputs["model"]

    # Gather BQML evaluation metrics using a pre-built-component.
    bqml_evaluate_op = bq_components.BigqueryEvaluateModelJobOp(
        project=project, location=bq_location, model=bqml_model
    )
    bqml_eval_metrics_raw = bqml_evaluate_op.outputs["evaluation_metrics"]

    # Analyzes evaluation BQML metrics using a custom component.
    interpret_bqml_evaluation_metrics_op = interpret_bqml_evaluation_metrics(
        bqml_evaluation_metrics=bqml_eval_metrics_raw
    )
    bqml_eval_metrics = interpret_bqml_evaluation_metrics_op.outputs["metrics"]

    # Exports the BQML model to a GCS bucket using a pre-built-component.
    bqml_export_op = bq_components.BigqueryExportModelJobOp(
        project=project,
        location=bq_location,
        model=bqml_model,
        model_destination_path=bqml_model_export_location,
    ).after(bqml_evaluate_op)
    bqml_exported_gcs_path = bqml_export_op.outputs["exported_model_path"]

    # Uploads the recently exported BQML model from GCS into Vertex AI using a pre-built-component.
    bqml_model_upload_op = vertex_pipeline_components.ModelUploadOp(
        project=project,
        location=region,
        display_name=DISPLAY_NAME + "_bqml",
        artifact_uri=bqml_exported_gcs_path,
        serving_container_image_uri=bqml_serving_container_image_uri,
    )
    bqml_vertex_model = bqml_model_upload_op.outputs["model"]

    # Creates a Vertex AI Tabular dataset using a pre-built-component.
    dataset_create_op = vertex_pipeline_components.TabularDatasetCreateOp(
        project=project,
        location=region,
        display_name=DISPLAY_NAME,
        bq_source=split_datasets_op.outputs["dataset_bq_uri"],
    )

    # Trains an AutoML Tables model using a pre-built-component.
    automl_training_op = vertex_pipeline_components.AutoMLTabularTrainingJobRunOp(
        project=project,
        location=region,
        display_name=f"{DISPLAY_NAME}_automl",
        optimization_prediction_type="regression",
        optimization_objective="minimize-rmse",
        predefined_split_column_name="split_col",
        dataset=dataset_create_op.outputs["dataset"],
        target_column="Rings",
        column_transformations=[
            {"categorical": {"column_name": "Sex"}},
            {"numeric": {"column_name": "Length"}},
            {"numeric": {"column_name": "Diameter"}},
            {"numeric": {"column_name": "Height"}},
            {"numeric": {"column_name": "Whole_weight"}},
            {"numeric": {"column_name": "Shucked_weight"}},
            {"numeric": {"column_name": "Viscera_weight"}},
            {"numeric": {"column_name": "Shell_weight"}},
            {"numeric": {"column_name": "Rings"}},
        ],
    )
    automl_model = automl_training_op.outputs["model"]

    # Analyzes evaluation AutoML metrics using a custom component.
    automl_eval_op = interpret_automl_evaluation_metrics(
        region=region, model=automl_model
    )
    automl_eval_metrics = automl_eval_op.outputs["metrics"]

    # 1) Decides which model is best (AutoML vs BQML);
    # 2) Determines if the best model meets the deployment condition.
    best_model_task = select_best_model(
        metrics_bqml=bqml_eval_metrics,
        metrics_automl=automl_eval_metrics,
        thresholds_dict_str=thresholds_dict_str,
    )

    # If the deploy condition is True, then deploy the best model.
    with dsl.Condition(
        best_model_task.outputs["deploy_decision"] == "true",
        name="deploy_decision",
    ):
        # Creates a Vertex AI endpoint using a pre-built-component.
        endpoint_create_op = vertex_pipeline_components.EndpointCreateOp(
            project=project,
            location=region,
            display_name=endpoint_display_name,
        )
        endpoint_create_op.after(best_model_task)

        # In case the BQML model is the best...
        with dsl.Condition(
            best_model_task.outputs["best_model"] == "bqml",
            name="deploy_bqml",
        ):
            # Deploys the BQML model (now on Vertex AI) to the recently created endpoint using a pre-built component.
            model_deploy_bqml_op = (
                vertex_pipeline_components.ModelDeployOp(  # noqa: F841
                    endpoint=endpoint_create_op.outputs["endpoint"],
                    model=bqml_vertex_model,
                    deployed_model_display_name=DISPLAY_NAME + "_best_bqml",
                    dedicated_resources_machine_type="n1-standard-2",
                    dedicated_resources_min_replica_count=2,
                    dedicated_resources_max_replica_count=2,
                    traffic_split={
                        "0": 100
                    },  # newly deployed model gets 100% of the traffic
                ).set_caching_options(False)
            )

            # Sends an online prediction request to the recently deployed model using a custom component.
            validate_infrastructure(
                endpoint=endpoint_create_op.outputs["endpoint"]
            ).set_caching_options(False).after(model_deploy_bqml_op)

        # In case the AutoML model is the best...
        with dsl.Condition(
            best_model_task.outputs["best_model"] == "automl",
            name="deploy_automl",
        ):
            # Deploys the AutoML model to the recently created endpoint using a pre-built component.
            model_deploy_automl_op = (
                vertex_pipeline_components.ModelDeployOp(  # noqa: F841
                    endpoint=endpoint_create_op.outputs["endpoint"],
                    model=automl_model,
                    deployed_model_display_name=DISPLAY_NAME + "_best_automl",
                    dedicated_resources_machine_type="n1-standard-2",
                    dedicated_resources_min_replica_count=2,
                    dedicated_resources_max_replica_count=2,
                    traffic_split={
                        "0": 100
                    },  # newly deployed model gets 100% of the traffic
                ).set_caching_options(False)
            )

            # Sends an online prediction request to the recently deployed model using a custom component.
            validate_infrastructure(
                endpoint=endpoint_create_op.outputs["endpoint"]
            ).set_caching_options(False).after(model_deploy_automl_op)

"""
Explanation: The Pipeline
End of explanation
"""


compiler.Compiler().compile(
    pipeline_func=train_pipeline,
    package_path=PIPELINE_JSON_PKG_PATH,
)


vertex.init(project=PROJECT_ID, location=REGION)

pipeline_job = # TODO 5: Your code goes here(
    display_name=DISPLAY_NAME,
    template_path=PIPELINE_JSON_PKG_PATH,
    pipeline_root=PIPELINE_ROOT,
    parameter_values=pipeline_params,
    enable_caching=False,
)

response = pipeline_job.submit()

"""
Explanation: Running the Pipeline
End of explanation
"""


pipeline_job.wait()

"""
Explanation: Wait for the pipeline to complete
Currently, your pipeline is running asynchronous by using the submit() method. To have run it synchronously, you would have invoked the run() method.
In this last step, you block on the asynchronously executed waiting for completion using the wait() method.
End of explanation
"""


vertex.init(project=PROJECT_ID, location=REGION)

delete_bucket = False

print("Will delete endpoint")
endpoints = vertex.Endpoint.list(
    filter=f"display_name={DISPLAY_NAME}_endpoint", order_by="create_time"
)
endpoint = endpoints[0]
endpoint.undeploy_all()
vertex.Endpoint.delete(endpoint)
print("Deleted endpoint:", endpoint)

print("Will delete models")
suffix_list = ["bqml", "automl", "best"]
for suffix in suffix_list:
    try:
        model_display_name = f"{DISPLAY_NAME}_{suffix}"
        print("Will delete model with name " + model_display_name)
        models = vertex.Model.list(
            filter=f"display_name={model_display_name}", order_by="create_time"
        )

        model = models[0]
        vertex.Model.delete(model)
        print("Deleted model:", model)
    except Exception as e:
        print(e)


print("Will delete Vertex dataset")
datasets = vertex.TabularDataset.list(
    filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
)

dataset = datasets[0]
vertex.TabularDataset.delete(dataset)
print("Deleted Vertex dataset:", dataset)


pipelines = vertex.PipelineJob.list(
    filter=f"pipeline_name={DISPLAY_NAME}", order_by="create_time"
)
pipeline = pipelines[0]
vertex.PipelineJob.delete(pipeline)
print("Deleted pipeline:", pipeline)

# Construct a BigQuery client object.

bq_client = bigquery.Client(project=PROJECT_ID, location=BQ_LOCATION)

# TODO(developer): Set dataset_id to the ID of the dataset to fetch.
dataset_id = f"{PROJECT_ID}.{BQ_DATASET}"

print(f"Will delete BQ dataset '{dataset_id}' from location {BQ_LOCATION}.")
# Use the delete_contents parameter to delete a dataset and its contents.
# Use the not_found_ok parameter to not receive an error if the dataset has already been deleted.
bq_client.delete_dataset(
    dataset_id, delete_contents=True, not_found_ok=True
)  # Make an API request.

print(f"Deleted BQ dataset '{dataset_id}' from location {BQ_LOCATION}.")

if delete_bucket or os.getenv("IS_TESTING"):
    ! gsutil rm -r $BUCKET_URI

"""
Explanation: Cleaning up
To clean up all Google Cloud resources used in this project, you can delete the Google Cloud
project you used for the notebook.
Otherwise, you can delete the individual resources you created in this notebook:
End of explanation
"""


