import os
import shutil
from datetime import datetime

import pandas as pd
import tensorflow as tf
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__)
%matplotlib inline

# Change below if necessary
PROJECT = !gcloud config get-value project  # noqa: E999
PROJECT = PROJECT[0]
BUCKET = PROJECT
REGION = "us-central1"

%env PROJECT=$PROJECT
%env BUCKET=$BUCKET
%env REGION=$REGION

%%bash
gcloud config set project $PROJECT
gcloud config set ai_platform/region $REGION

"""
Explanation: Training a model with traffic_last_5min feature
Introduction
In this notebook, we'll train a taxifare prediction model but this time with an additional feature of traffic_last_5min.
End of explanation
"""


!ls -l ../data/taxi-traffic*

!head ../data/taxi-traffic*

"""
Explanation: Load raw data
End of explanation
"""


CSV_COLUMNS = [
    "fare_amount",
    "dayofweek",
    "hourofday",
    "pickup_longitude",
    "pickup_latitude",
    "dropoff_longitude",
    "dropoff_latitude",
    "traffic_last_5min",
]
LABEL_COLUMN = "fare_amount"
DEFAULTS = [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]]


def features_and_labels(row_data):
    label = row_data.pop(LABEL_COLUMN)
    features = row_data

    return features, label


def create_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):
    dataset = tf.data.experimental.make_csv_dataset(
        pattern, batch_size, CSV_COLUMNS, DEFAULTS
    )

    dataset = dataset.map(features_and_labels)

    if mode == tf.estimator.ModeKeys.TRAIN:
        dataset = dataset.shuffle(buffer_size=1000).repeat()

    # take advantage of multi-threading; 1=AUTOTUNE
    dataset = dataset.prefetch(1)
    return dataset

INPUT_COLS = [
    "dayofweek",
    "hourofday",
    "pickup_longitude",
    "pickup_latitude",
    "dropoff_longitude",
    "dropoff_latitude",
    "traffic_last_5min",
]

# Create input layer of feature columns
feature_columns = {
    colname: tf.feature_column.numeric_column(colname) for colname in INPUT_COLS
}

"""
Explanation: Use tf.data to read the CSV files
These functions for reading data from the csv files are similar to what we used in the Introduction to Tensorflow module. Note that here we have an addtional feature traffic_last_5min.
End of explanation
"""


# Build a keras DNN model using Sequential API
def build_model(dnn_hidden_units):
    model = Sequential(DenseFeatures(feature_columns=feature_columns.values()))

    for num_nodes in dnn_hidden_units:
        model.add(Dense(units=num_nodes, activation="relu"))

    model.add(Dense(units=1, activation="linear"))

    # Create a custom evaluation metric
    def rmse(y_true, y_pred):
        return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))

    # Compile the keras model
    model.compile(optimizer="adam", loss="mse", metrics=[rmse, "mse"])

    return model

"""
Explanation: Build a simple keras DNN model
End of explanation
"""


HIDDEN_UNITS = [32, 8]

model = build_model(dnn_hidden_units=HIDDEN_UNITS)

BATCH_SIZE = 1000
NUM_TRAIN_EXAMPLES = 10000 * 6  # training dataset will repeat, wrap around
NUM_EVALS = 60  # how many times to evaluate
NUM_EVAL_EXAMPLES = 10000  # enough to get a reasonable sample

trainds = create_dataset(
    pattern="../data/taxi-traffic-train*",
    batch_size=BATCH_SIZE,
    mode=tf.estimator.ModeKeys.TRAIN,
)

evalds = create_dataset(
    pattern="../data/taxi-traffic-valid*",
    batch_size=BATCH_SIZE,
    mode=tf.estimator.ModeKeys.EVAL,
).take(NUM_EVAL_EXAMPLES // 1000)

%%time
steps_per_epoch = NUM_TRAIN_EXAMPLES // (BATCH_SIZE * NUM_EVALS)

LOGDIR = "./taxi_trained"
history = model.fit(
    x=trainds,
    steps_per_epoch=steps_per_epoch,
    epochs=NUM_EVALS,
    validation_data=evalds,
    callbacks=[TensorBoard(LOGDIR)],
)

RMSE_COLS = ["rmse", "val_rmse"]

pd.DataFrame(history.history)[RMSE_COLS].plot()

model.predict(
    x={
        "dayofweek": tf.convert_to_tensor([6]),
        "hourofday": tf.convert_to_tensor([17]),
        "pickup_longitude": tf.convert_to_tensor([-73.982683]),
        "pickup_latitude": tf.convert_to_tensor([40.742104]),
        "dropoff_longitude": tf.convert_to_tensor([-73.983766]),
        "dropoff_latitude": tf.convert_to_tensor([40.755174]),
        "traffic_last_5min": tf.convert_to_tensor([114]),
    },
    steps=1,
)

"""
Explanation: Next, we can call the build_model to create the model. Here we'll have two hidden layers before our final output layer. And we'll train with the same parameters we used before.
End of explanation
"""


OUTPUT_DIR = "./export/savedmodel"
shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
EXPORT_PATH = os.path.join(OUTPUT_DIR, datetime.now().strftime("%Y%m%d%H%M%S"))
model.save(EXPORT_PATH)  # with default serving function
os.environ["EXPORT_PATH"] = EXPORT_PATH

%%bash
PROJECT=${PROJECT}
BUCKET=${BUCKET}
REGION=global
MODEL_NAME=taxifare
VERSION_NAME=traffic

if [[ $(gcloud ai-platform models list --format='value(name)' --region=$REGION | grep "^$MODEL_NAME$") ]]; then
    echo "$MODEL_NAME already exists"
else
    # create model
    echo "Creating $MODEL_NAME"
    gcloud ai-platform models create --region=$REGION $MODEL_NAME
fi

if [[ $(gcloud ai-platform versions list --model $MODEL_NAME --format='value(name)' --region=$REGION | grep "^$VERSION_NAME$") ]]; then
    echo "Deleting already existing $MODEL_NAME:$VERSION_NAME ... "
    gcloud ai-platform versions delete --model=$MODEL_NAME $VERSION_NAME --region=$REGION
    echo "Please run this cell again if you don't see a Creating message ... "
    sleep 2
fi

# create model
echo "Creating $MODEL_NAME:$VERSION_NAME"
gcloud ai-platform versions create --model=$MODEL_NAME $VERSION_NAME --async \
       --framework=tensorflow --python-version=3.7 --runtime-version=2.5 \
       --origin=${EXPORT_PATH} --staging-bucket=gs://$BUCKET --region=$REGION

"""
Explanation: Export and deploy model
End of explanation
"""


