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"

! pip uninstall {USER_FLAG} -y google-cloud-aiplatform
! pip uninstall {USER_FLAG} -y google-cloud-bigquery
! pip uninstall {USER_FLAG} -y google-cloud-bigquery-storage
! pip uninstall {USER_FLAG} -y google-cloud-aiplatform

! pip install {USER_FLAG} --upgrade google-cloud-aiplatform
! pip install {USER_FLAG} --upgrade google-cloud-bigquery
! pip install {USER_FLAG} --upgrade google-cloud-bigquery-storage
! pip install {USER_FLAG} avro

"""
Explanation: <table align="left">
  <td>
    <a href="https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/sdk-feature-store-pandas.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/colab-logo-32px.png" alt="Colab logo"> Run in Colab
    </a>
  </td>
  <td>
    <a href="https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/sdk-feature-store-pandas.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/github-logo-32px.png" alt="GitHub logo">
      View on GitHub
    </a>
  </td>
</table>

Overview
This Colab introduces Pandas support of Vertex AI SDK Feature Store. For pre-requisite and introduction for Vertex AI SDK Feature Store native support, please see this Colab. 
Dataset
This Colab uses a movie recommendation dataset as an example throughout all the sessions. The task is to train a model to predict if a user is going to watch a movie and serve this model online.
Objective
In this notebook, you will learn how to:
* Ingest Feature Values from Pandas DataFrame into featurestore's entity types.
* Read Entity Feature Values from Online Feature Store into Pandas DataFrame.
* Batch Serve Feature Values from your featurestore to Pandas DataFrame.

We will also discuss how Vertex AI Feature Store can be useful in the below scenarios:
* online serving with updated feature values
* point-in-time correctness to fetch feature values for training

Costs
This tutorial uses billable components of Google Cloud:

Vertex AI
Cloud BigQuery

Learn about Vertex AI
pricing and Cloud Storage
pricing, and use the Pricing
Calculator
to generate a cost estimate based on your projected usage.
Before you begin
Install additional packages
For this Colab, you need the Vertex SDK for Python.
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: Restart the kernel
After you install the SDK, you need to restart the notebook kernel so it can find the packages. You can restart kernel from Kernel -> Restart Kernel, or running the following:
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 up your Google Cloud project
The following steps are required, regardless of your notebook environment.


Select or create a Google Cloud project. When you first create an account, you get a $300 free credit towards your compute/storage costs.


Make sure that billing is enabled for your project.


Enable the Vertex AI API and Compute Engine API.


If you are running this notebook locally, you will need to install the Cloud SDK.


Enter your project ID in the cell below. Then run the cell to make sure the
Cloud SDK uses the right project for all the commands in this notebook.


Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands.
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 = ""  # @param {type:"string"}
print("Project ID: ", PROJECT_ID)

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


import os
import sys

# If you are running this notebook in Colab, run this cell and follow the
# instructions to authenticate your GCP account. This provides access to your
# Cloud Storage bucket and lets you submit training jobs and prediction
# requests.

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

# If on Google Cloud Notebooks, then don't execute this code
if not IS_GOOGLE_CLOUD_NOTEBOOK:
    if "google.colab" in sys.modules:
        from google.colab import auth as google_auth

        google_auth.authenticate_user()

    # If you are running this notebook locally, replace the string below with the
    # path to your service account key and run this cell to authenticate your GCP
    # account.
    elif not os.getenv("IS_TESTING"):
        %env GOOGLE_APPLICATION_CREDENTIALS ''

"""
Explanation: Authenticate your Google Cloud account
If you are using Google Cloud Notebooks, your environment is already
authenticated. Skip this step.
If you are using Colab, run the cell below and follow the instructions
when prompted to authenticate your account via oAuth.
Otherwise, follow these steps:


In the Cloud Console, go to the Create service account key
   page.


Click Create service account.


In the Service account name field, enter a name, and
   click Create.


In the Grant this service account access to project section, click the Role drop-down list. Type "Vertex AI"
into the filter box, and select
   Vertex AI Administrator. Type "Storage Object Admin" into the filter box, and select Storage Object Admin.


Click Create. A JSON file that contains your key downloads to your
local environment.


Enter the path to your service account key as the
GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell.
End of explanation
"""


import datetime

import pandas as pd
from google.cloud import aiplatform

REGION = ""  # @param {type:"string"}

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

"""
Explanation: Import libraries and define constants
End of explanation
"""


movie_predictions_feature_store = aiplatform.Featurestore.create(
    featurestore_id="movie_predictions",
    online_store_fixed_node_count=1,
)

"""
Explanation: Create Feature Store Resources
Create Featurestore
The method to create a Featurestore returns a
long-running operation (LRO). An LRO starts an asynchronous job. LROs are returned for other API
methods too, such as updating or deleting a featurestore. Running the code cell will create a featurestore and print the process log.
End of explanation
"""


users_entity_type = movie_predictions_feature_store.create_entity_type(
    entity_type_id="users",
    description="Users entity",
)

movies_entity_type = movie_predictions_feature_store.create_entity_type(
    entity_type_id="movies",
    description="Movies entity",
)

"""
Explanation: Create Entity Types
Entity types can be created within the Featurestore class. Below, create the Users entity type and Movies entity type. A process log will be printed out.
End of explanation
"""


users_feature_age = users_entity_type.create_feature(
    feature_id="age",
    value_type="INT64",
    description="User age",
)

users_feature_gender = users_entity_type.create_feature(
    feature_id="gender",
    value_type="STRING",
    description="User gender",
)

users_feature_liked_genres = users_entity_type.create_feature(
    feature_id="liked_genres",
    value_type="STRING_ARRAY",
    description="An array of genres this user liked",
)

movies_feature_configs = {
    "title": {
        "value_type": "STRING",
        "description": "The title of the movie",
    },
    "genres": {
        "value_type": "STRING",
        "description": "The genre of the movie",
    },
    "average_rating": {
        "value_type": "DOUBLE",
        "description": "The average rating for the movie, range is [1.0-5.0]",
    },
}

movie_features = movies_entity_type.batch_create_features(
    feature_configs=movies_feature_configs,
)

"""
Explanation: Create Features
Features can be created within each entity type. Add defining features to the Users entity type and Movies entity type by using the following methods.
End of explanation
"""


GCS_USERS_AVRO_URI = (
    "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/users.avro"
)
GCS_MOVIES_AVRO_URI = (
    "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movies.avro"
)

USERS_AVRO_FN = "users.avro"
MOVIES_AVRO_FN = "movies.avro"

! gsutil cp $GCS_USERS_AVRO_URI $USERS_AVRO_FN
! gsutil cp $GCS_MOVIES_AVRO_URI $MOVIES_AVRO_FN

"""
Explanation: Ingest Feature Values into Entity Type from a Pandas DataFrame
You need to ingest feature values into your entity type containing the features, so you can later read (online) or batch serve (offline) the feature values from the entity type. In this step, you will learn how to ingest feature values from a Pandas DataFrame into an entity type. We can also import feature values from BigQuery or Google Cloud Storage.
Entity Type Source Files
End of explanation
"""


from avro.datafile import DataFileReader
from avro.io import DatumReader


class AvroReader:
    def __init__(self, data_file):
        self.avro_reader = DataFileReader(open(data_file, "rb"), DatumReader())

    def to_dataframe(self):
        records = [record for record in self.avro_reader]
        return pd.DataFrame.from_records(data=records)

users_avro_reader = AvroReader(data_file=USERS_AVRO_FN)
users_source_df = users_avro_reader.to_dataframe()
print(users_source_df)

movies_avro_reader = AvroReader(data_file=MOVIES_AVRO_FN)
movies_source_df = movies_avro_reader.to_dataframe()
print(movies_source_df)

"""
Explanation: Load Avro Files into Pandas DataFrames
End of explanation
"""


users_entity_type.ingest_from_df(
    feature_ids=["age", "gender", "liked_genres"],
    feature_time="update_time",
    df_source=users_source_df,
    entity_id_field="user_id",
)

"""
Explanation: Ingest Feature Values into Users Entity Type
End of explanation
"""


movies_entity_type.ingest_from_df(
    feature_ids=["average_rating", "title", "genres"],
    feature_time="update_time",
    df_source=movies_source_df,
    entity_id_field="movie_id",
)

"""
Explanation: Ingest Feature Values into Movies Entity Type
End of explanation
"""


users_read_df = users_entity_type.read(
    entity_ids=["dave", "alice", "charlie", "bob", "eve"],
)
print(users_read_df)

movies_read_df = movies_entity_type.read(
    entity_ids=["movie_01", "movie_02", "movie_03", "movie_04"],
    feature_ids=["title", "genres", "average_rating"],
)
print(movies_read_df)

"""
Explanation: Read/Online Serve Entity's Feature Values from Vertex AI Online Feature Store
Feature Store allows online serving
which lets you read feature values for small batches of entities. It works well when you want to read values of selected features from an entity or multiple entities in an entity type.
End of explanation
"""


GCS_READ_INSTANCES_CSV_URI = "gs://cloud-samples-data-us-central1/vertex-ai/feature-store/datasets/movie_prediction.csv"

! gsutil cp $GCS_READ_INSTANCES_CSV_URI $READ_INSTANCES_CSV_FN

"""
Explanation: Batch Serve Featurestore's Feature Values from Vertex AI Feature Store
Batch Serving is used to fetch a large batch of feature values for high-throughput, and is typically used for training a model or batch prediction. In this section, you will learn how to prepare for training examples by using the Featurestore's batch serve function.
Read Instances Source File
End of explanation
"""


read_instances_df = pd.read_csv(read_instances_csv_fn)
print(read_instances_df)

"""
Explanation: Load Csv File into a Pandas DataFrame
End of explanation
"""


print("before: ", read_instances_df["timestamp"].dtype)
read_instances_df = read_instances_df.astype({"timestamp": "datetime64"})
print("after:  ", read_instances_df["timestamp"].dtype)

"""
Explanation: Change the Dtype of Timestamp to Datetime64
End of explanation
"""


movie_predictions_df = movie_predictions_feature_store.batch_serve_to_df(
    serving_feature_ids={
        "users": ["age", "gender", "liked_genres"],
        "movies": ["title", "average_rating", "genres"],
    },
    read_instances_df=read_instances_df,
)
movie_predictions_df

"""
Explanation: Batch Serve Feature Values from Movie Predictions Featurestore
End of explanation
"""


print(movies_read_df)

"""
Explanation: Read the Updated Feature Values
Recall Read from the Entity Type Shows Feature Values from the Last Ingestion
End of explanation
"""


update_movies_df = pd.DataFrame(
    data=[["movie_03", 4.3], ["movie_04", 4.8]],
    columns=["movie_id", "average_rating"],
)
print(update_movies_df)

movies_entity_type.ingest_from_df(
    feature_ids=["average_rating"],
    feature_time=datetime.datetime.now(),
    df_source=update_movies_df,
    entity_id_field="movie_id",
)

"""
Explanation: Ingest Updated Feature Values
End of explanation
"""


update_movies_read_df = movies_entity_type.read(
    entity_ids=["movie_01", "movie_02", "movie_03", "movie_04"],
    feature_ids=["title", "genres", "average_rating"],
)
print(update_movies_read_df)

"""
Explanation: Read from the Entity Type Shows Updated Feature Values from the Latest Ingestion
End of explanation
"""


print(movie_predictions_df)

"""
Explanation: Point-in-Time Correctness
Recall Batch Serve From the Last Ingestion Has Missing Data
End of explanation
"""


backfill_users_df = pd.DataFrame(
    data=[["bob", 34, "Male", ["Drama"], "2020-02-13 09:35:15"]],
    columns=["user_id", "age", "gender", "liked_genres", "update_time"],
)
backfill_users_df = backfill_users_df.astype({"update_time": "datetime64"})
print(backfill_users_df)

backfill_movies_df = pd.DataFrame(
    data=[["movie_04", 4.2, "The Dark Knight", "Action", "2020-02-13 09:35:15"]],
    columns=["movie_id", "average_rating", "title", "genres", "update_time"],
)
backfill_movies_df = backfill_movies_df.astype({"update_time": "datetime64"})
print(backfill_movies_df)

"""
Explanation: Backfill/Correct Point-in-Time Data
End of explanation
"""


users_entity_type.ingest_from_df(
    feature_ids=["age", "gender", "liked_genres"],
    feature_time="update_time",
    df_source=backfill_users_df,
    entity_id_field="user_id",
)

movies_entity_type.ingest_from_df(
    feature_ids=["average_rating", "title", "genres"],
    feature_time="update_time",
    df_source=backfill_movies_df,
    entity_id_field="movie_id",
)

"""
Explanation: Ingest Backfill/Correct Point-in-Time Data
End of explanation
"""


backfill_movie_predictions_df = movie_predictions_feature_store.batch_serve_to_df(
    serving_feature_ids={
        "users": ["age", "gender", "liked_genres"],
        "movies": ["title", "average_rating", "genres"],
    },
    read_instances_df=read_instances_df,
)
print(backfill_movie_predictions_df)

"""
Explanation: Batch Serve From the Latest Ingestion with Backfill/Correction Has Reduced Missing Data
End of explanation
"""


movie_predictions_feature_store.delete(force=True)

"""
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 tutorial.
You can also keep the project but delete the featurestore by running the code below:
End of explanation
"""


