import numpy as np
import os
import shutil
import tensorflow as tf

from google.cloud import aiplatform
from google.cloud import bigquery
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Value
from matplotlib import pyplot as plt
from tensorflow import keras
from tensorflow.keras.callbacks import TensorBoard
from tensorflow.keras.layers import Dense, DenseFeatures
from tensorflow.keras.models import Sequential

print(tf.__version__)

PROJECT = 'cloud-training-demos' # REPLACE WITH YOUR PROJECT ID
BUCKET = 'cloud-training-demos' # REPLACE WITH YOUR BUCKET NAME
REGION = 'us-central1' # REPLACE WITH YOUR BUCKET REGION e.g. us-central1

# For Bash Code
os.environ['PROJECT'] = PROJECT
os.environ['BUCKET'] = BUCKET
os.environ['REGION'] = REGION

%%bash
gcloud config set project $PROJECT

"""
Explanation: Working with Streaming Data
Learning Objectives
 1. Learn how to process real-time data for ML models using Cloud Dataflow
 2. Learn how to serve online predictions using real-time data
Introduction
It can be useful to leverage real time data in a machine learning model when making a prediction. However, doing so requires setting up a streaming data pipeline which can be non-trivial. 
Typically you will have the following:
 - A series of IoT devices generating and sending data from the field in real-time (in our case these are the taxis)
 - A messaging bus to that receives and temporarily stores the IoT data (in our case this is Cloud Pub/Sub)
 - A streaming processing service that subscribes to the messaging bus, windows the messages and performs data transformations on each window (in our case this is Cloud Dataflow)
 - A persistent store to keep the processed data (in our case this is BigQuery)
These steps happen continuously and in real-time, and are illustrated by the blue arrows in the diagram below. 
Once this streaming data pipeline is established, we need to modify our model serving to leverage it. This simply means adding a call to the persistent store (BigQuery) to fetch the latest real-time data when a prediction request comes in. This flow is illustrated by the red arrows in the diagram below. 
<img src='../assets/taxi_streaming_data.png' width='80%'>
In this lab we will address how to process real-time data for machine learning models. We will use the same data as our previous 'taxifare' labs, but with the addition of trips_last_5min data as an additional feature. This is our proxy for real-time traffic.
End of explanation
"""


bq = bigquery.Client()

dataset = bigquery.Dataset(bq.dataset("taxifare"))
try:
    bq.create_dataset(dataset)  # will fail if dataset already exists
    print("Dataset created.")
except:
    print("Dataset already exists.")

"""
Explanation: Re-train our model with trips_last_5min feature
In this lab, we want to show how to process real-time data for training and prediction. So, we need to retrain our previous model with this additional feature. Go through the notebook 4a_streaming_data_training.ipynb. Open and run the notebook to train and save a model. This notebook is very similar to what we did in the Introduction to Tensorflow module but note the added feature for trips_last_5min in the model and the dataset.
Simulate Real Time Taxi Data
Since we don’t actually have real-time taxi data we will synthesize it using a simple python script. The script publishes events to Google Cloud Pub/Sub.
Inspect the iot_devices.py script in the taxicab_traffic folder. It is configured to send about 2,000 trip messages every five minutes with some randomness in the frequency to mimic traffic fluctuations. These numbers come from looking at the historical average of taxi ride frequency in BigQuery. 
In production this script would be replaced with actual taxis with IoT devices sending trip data to Cloud Pub/Sub. 
To execute the iot_devices.py script, launch a terminal and navigate to the asl-ml-immersion/notebooks/building_production_ml_systems/solutions directory. Then run the following two commands.
bash
PROJECT_ID=$(gcloud config list project --format "value(core.project)")
python3 ./taxicab_traffic/iot_devices.py --project=$PROJECT_ID
You will see new messages being published every 5 seconds. Keep this terminal open so it continues to publish events to the Pub/Sub topic. If you open Pub/Sub in your Google Cloud Console, you should be able to see a topic called taxi_rides.
Create a BigQuery table to collect the processed data
In the next section, we will create a dataflow pipeline to write processed taxifare data to a BigQuery Table, however that table does not yet exist. Execute the following commands to create a BigQuery dataset called taxifare and a table within that dataset called traffic_realtime.
End of explanation
"""


dataset = bigquery.Dataset(bq.dataset("taxifare"))

table_ref = dataset.table("traffic_realtime")
SCHEMA = [
    bigquery.SchemaField("trips_last_5min", "INTEGER", mode="REQUIRED"),
    bigquery.SchemaField("time", "TIMESTAMP", mode="REQUIRED"),
]
table = bigquery.Table(table_ref, schema=SCHEMA)

try:
    bq.create_table(table)
    print("Table created.")
except:
    print("Table already exists.")

"""
Explanation: Next, we create a table called traffic_realtime and set up the schema.
End of explanation
"""


%%bigquery
SELECT
  *
FROM
  `taxifare.traffic_realtime`
ORDER BY
  time DESC
LIMIT 10

"""
Explanation: Launch Streaming Dataflow Pipeline
Now that we have our taxi data being pushed to Pub/Sub, and our BigQuery table set up, let’s consume the Pub/Sub data using a streaming DataFlow pipeline.
The pipeline is defined in ./taxicab_traffic/streaming_count.py. Open that file and inspect it. 
There are 5 transformations being applied:
 - Read from PubSub
 - Window the messages
 - Count number of messages in the window
 - Format the count for BigQuery
 - Write results to BigQuery
TODO: Open the file ./taxicab_traffic/streaming_count.py and find the TODO there. Specify a sliding window that is 5 minutes long, and gets recalculated every 15 seconds. Hint: Reference the beam programming guide for guidance. To check your answer reference the solution. 
For the second transform, we specify a sliding window that is 5 minutes long, and recalculate values every 15 seconds. 
In a new terminal, launch the dataflow pipeline using the command below. You can change the BUCKET variable, if necessary. Here it is assumed to be your PROJECT_ID.
bash
PROJECT_ID=$(gcloud config list project --format "value(core.project)")
BUCKET=$PROJECT_ID # CHANGE AS NECESSARY 
python3 ./taxicab_traffic/streaming_count.py \
    --input_topic taxi_rides \
    --runner=DataflowRunner \
    --project=$PROJECT_ID \
    --temp_location=gs://$BUCKET/dataflow_streaming
Once you've submitted the command above you can examine the progress of that job in the Dataflow section of Cloud console. 
Explore the data in the table
After a few moments, you should also see new data written to your BigQuery table as well. 
Re-run the query periodically to observe new data streaming in! You should see a new row every 15 seconds.
End of explanation
"""


# TODO 2a. Write a function to take most recent entry in `traffic_realtime` table and add it to instance.
def add_traffic_last_5min(instance):
    bq = bigquery.Client()
    query_string = """
    SELECT
      *
    FROM
      `taxifare.traffic_realtime`
    ORDER BY
      time DESC
    LIMIT 1
    """
    trips = bq.query(query_string).to_dataframe()['trips_last_5min'][0]
    instance['traffic_last_5min'] = int(trips)
    return instance

"""
Explanation: Make predictions from the new data
In the rest of the lab, we'll referece the model we trained and deployed from the previous labs, so make sure you have run the code in the 4a_streaming_data_training.ipynb notebook. 
The add_traffic_last_5min function below will query the traffic_realtime table to find the most recent traffic information and add that feature to our instance for prediction.
End of explanation
"""


add_traffic_last_5min(instance={'dayofweek': 4,
                                'hourofday': 13,
                                'pickup_longitude': -73.99,
                                'pickup_latitude': 40.758,
                                'dropoff_latitude': 41.742,
                                'dropoff_longitude': -73.07})

"""
Explanation: The traffic_realtime table is updated in realtime using Cloud Pub/Sub and Dataflow so, if you run the cell below periodically, you should see the traffic_last_5min feature added to the instance and change over time.
End of explanation
"""


# TODO 2b. Write code to call prediction on instance using realtime traffic info.
#Hint: Look at this sample https://github.com/googleapis/python-aiplatform/blob/master/samples/snippets/predict_custom_trained_model_sample.py

ENDPOINT_ID = # TODO: Copy the `ENDPOINT_ID` from the deployment in the previous lab.

api_endpoint = f'{REGION}-aiplatform.googleapis.com'

# The AI Platform services require regional API endpoints.
client_options = {"api_endpoint": api_endpoint}
# Initialize client that will be used to create and send requests.
# This client only needs to be created once, and can be reused for multiple requests.
client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)

instance = {'dayofweek': 4,
            'hourofday': 13,
            'pickup_longitude': -73.99,
            'pickup_latitude': 40.758,
            'dropoff_latitude': 41.742,
            'dropoff_longitude': -73.07}

# The format of each instance should conform to the deployed model's prediction input schema.
instance_dict = add_traffic_last_5min(instance)

instance = json_format.ParseDict(instance, Value())
instances = [instance]
endpoint = client.endpoint_path(
    project=PROJECT, location=REGION, endpoint=ENDPOINT_ID
)
response = client.predict(
    endpoint=endpoint, instances=instances
)

# The predictions are a google.protobuf.Value representation of the model's predictions.
print(" prediction:", response.predictions[0][0])


"""
Explanation: Finally, we'll use the python api to call predictions on an instance, using the realtime traffic information in our prediction. Just as above, you should notice that our resulting predicitons change with time as our realtime traffic information changes as well.
Copy the ENDPOINT_ID from the deployment in the previous lab to the beginning of the block below.
End of explanation
"""


