import os
import tempfile
import time
from pprint import pprint

import absl
import pandas as pd
import tensorflow as tf
import tensorflow_data_validation as tfdv
import tensorflow_model_analysis as tfma
import tensorflow_transform as tft
import tfx
from tensorflow_metadata.proto.v0 import (
    anomalies_pb2,
    schema_pb2,
    statistics_pb2,
)
from tensorflow_transform.tf_metadata import schema_utils
from tfx.components import (
    CsvExampleGen,
    Evaluator,
    ExampleValidator,
    InfraValidator,
    Pusher,
    ResolverNode,
    SchemaGen,
    StatisticsGen,
    Trainer,
    Transform,
    Tuner,
)
from tfx.components.common_nodes.importer_node import ImporterNode
from tfx.components.trainer import executor as trainer_executor
from tfx.dsl.components.base import executor_spec
from tfx.dsl.experimental import latest_blessed_model_resolver
from tfx.orchestration import metadata, pipeline
from tfx.orchestration.experimental.interactive.interactive_context import (
    InteractiveContext,
)
from tfx.proto import (
    evaluator_pb2,
    example_gen_pb2,
    infra_validator_pb2,
    pusher_pb2,
    trainer_pb2,
)
from tfx.proto.evaluator_pb2 import SingleSlicingSpec
from tfx.types import Channel
from tfx.types.standard_artifacts import (
    HyperParameters,
    InfraBlessing,
    Model,
    ModelBlessing,
)

"""
Explanation: Text Classification TFX Pipeline Starter
Objective: In this notebook, we show you how to put a text classification model implemented in model.py, preprocessing.py, and config.py into an interactive TFX pipeline. Using these files and the code snippets in this notebook, you'll configure a TFX pipeline generated by the tfx template tool as in the previous guided project so that the text classification can be run on a CAIP Pipelines Kubeflow cluster. The dataset itself consists of article titles along with their source, and the goal is to predict the source from the title. (This dataset can be re-generated by running either the keras_for_text_classification.ipynb notebook or the reusable_embeddings.ipynb notebook, which contain different models to solve this problem.) The solution we propose here is fairly simple and you can build on it by inspecting these notebooks.
End of explanation
"""


print("Tensorflow Version:", tf.__version__)
print("TFX Version:", tfx.__version__)
print("TFDV Version:", tfdv.__version__)
print("TFMA Version:", tfma.VERSION_STRING)

absl.logging.set_verbosity(absl.logging.INFO)

"""
Explanation: Note: this lab was developed and tested with the following TF ecosystem package versions:
Tensorflow Version: 2.3.1
TFX Version: 0.25.0
TFDV Version: 0.25.0
TFMA Version: 0.25.0
If you encounter errors with the above imports (e.g. TFX component not found), check your package versions in the cell below.
End of explanation
"""


os.environ["PATH"] += os.pathsep + "/home/jupyter/.local/bin"

"""
Explanation: If the versions above do not match, update your packages in the current Jupyter kernel below. The default %pip package installation location is not on your system installation PATH; use the command below to append the local installation path to pick up the latest package versions. Note that you may also need to restart your notebook kernel to pick up the specified package versions and re-run the imports cell above before proceeding with the lab.
End of explanation
"""


ARTIFACT_STORE = os.path.join(os.sep, "home", "jupyter", "artifact-store")
SERVING_MODEL_DIR = os.path.join(os.sep, "home", "jupyter", "serving_model")
DATA_ROOT = "./data"

DATA_ROOT = f"{ARTIFACT_STORE}/data"
!mkdir -p $DATA_ROOT

"""
Explanation: Configure lab settings
Set constants, location paths and other environment settings.
End of explanation
"""


data = pd.read_csv("./data/titles_sample.csv")
data.head()

LABEL_MAPPING = {"github": 0, "nytimes": 1, "techcrunch": 2}
data["source"] = data["source"].apply(lambda label: LABEL_MAPPING[label])
data.head()

data.to_csv(f"{DATA_ROOT}/dataset.csv", index=None)
!head $DATA_ROOT/*.csv

"""
Explanation: Preparing the dataset
End of explanation
"""


PIPELINE_NAME = "tfx-title-classifier"
PIPELINE_ROOT = os.path.join(
    ARTIFACT_STORE, PIPELINE_NAME, time.strftime("%Y%m%d_%H%M%S")
)
os.makedirs(PIPELINE_ROOT, exist_ok=True)

context = InteractiveContext(
    pipeline_name=PIPELINE_NAME,
    pipeline_root=PIPELINE_ROOT,
    metadata_connection_config=None,
)

"""
Explanation: Interactive Context
TFX Interactive Context allows you to create and run TFX Components in an interactive mode. It is designed to support experimentation and development in a Jupyter Notebook environment. It is an experimental feature and major changes to interface and functionality are expected. When creating the interactive context you can specifiy the following parameters:
- pipeline_name - Optional name of the pipeline for ML Metadata tracking purposes. If not specified, a name will be generated for you.
- pipeline_root - Optional path to the root of the pipeline's outputs. If not specified, an ephemeral temporary directory will be created and used.
- metadata_connection_config - Optional metadata_store_pb2.ConnectionConfig instance used to configure connection to a ML Metadata connection. If not specified, an ephemeral SQLite MLMD connection contained in the pipeline_root directory with file name "metadata.sqlite" will be used.
End of explanation
"""


output_config = example_gen_pb2.Output(
    split_config=example_gen_pb2.SplitConfig(
        splits=[
            example_gen_pb2.SplitConfig.Split(name="train", hash_buckets=4),
            example_gen_pb2.SplitConfig.Split(name="eval", hash_buckets=1),
        ]
    )
)

example_gen = tfx.components.CsvExampleGen(
    input_base=DATA_ROOT, output_config=output_config
)

context.run(example_gen)

"""
Explanation: Ingesting data using ExampleGen
In any ML development process the first step  is to ingest the training and test datasets. The ExampleGen component ingests data into a TFX pipeline. It consumes external files/services to generate a set file files in the TFRecord format,  which will be used by other TFX components. It can also shuffle the data and split into an arbitrary number of partitions.
<img src=../../images/ExampleGen.png width="300">
Configure and run CsvExampleGen
End of explanation
"""


examples_uri = example_gen.outputs["examples"].get()[0].uri

tfrecord_filenames = [
    os.path.join(examples_uri, "train", name)
    for name in os.listdir(os.path.join(examples_uri, "train"))
]

dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type="GZIP")

for tfrecord in dataset.take(2):
    example = tf.train.Example()
    example.ParseFromString(tfrecord.numpy())

    for name, feature in example.features.feature.items():
        if feature.HasField("bytes_list"):
            value = feature.bytes_list.value
        if feature.HasField("float_list"):
            value = feature.float_list.value
        if feature.HasField("int64_list"):
            value = feature.int64_list.value
        print(f"{name}: {value}")
    print("******")

"""
Explanation: Examine the ingested data
End of explanation
"""


statistics_gen = tfx.components.StatisticsGen(
    examples=example_gen.outputs["examples"]
)

context.run(statistics_gen)

"""
Explanation: Generating statistics using StatisticsGen
The StatisticsGen  component generates data statistics that can be used by other TFX components. StatisticsGen uses TensorFlow Data Validation. StatisticsGen generate statistics for each split in the ExampleGen component's output. In our case there two splits: train and eval.
<img src=../../images/StatisticsGen.png width="200">
Configure and  run the StatisticsGen component
End of explanation
"""


context.show(statistics_gen.outputs["statistics"])

"""
Explanation: Visualize statistics
The generated statistics can be visualized using the tfdv.visualize_statistics() function from the TensorFlow Data Validation library or using a utility method of the InteractiveContext object. In fact, most of the artifacts generated by the TFX components can be visualized using InteractiveContext.
End of explanation
"""


schema_gen = SchemaGen(
    statistics=statistics_gen.outputs["statistics"], infer_feature_shape=False
)

context.run(schema_gen)

"""
Explanation: Infering data schema using SchemaGen
Some TFX components use a description input data called a schema. The schema is an instance of schema.proto. It can specify data types for feature values, whether a feature has to be present in all examples, allowed value ranges, and other properties. SchemaGen automatically generates the schema by inferring types, categories, and ranges from data statistics. The auto-generated schema is best-effort and only tries to infer basic properties of the data. It is expected that developers review and modify it as needed. SchemaGen uses TensorFlow Data Validation.
The SchemaGen component generates the schema using the statistics for the train split. The statistics for other splits are ignored.
<img src=../../images/SchemaGen.png width="200">
Configure and run the SchemaGen components
End of explanation
"""


context.show(schema_gen.outputs["schema"])

"""
Explanation: Visualize the inferred schema
End of explanation
"""


schema_proto_path = "{}/{}".format(
    schema_gen.outputs["schema"].get()[0].uri, "schema.pbtxt"
)
schema = tfdv.load_schema_text(schema_proto_path)

"""
Explanation: Updating the auto-generated schema
In most cases the auto-generated schemas must be fine-tuned manually using insights from data exploration and/or domain knowledge about the data. For example, you know that in the covertype dataset there are seven types of forest cover (coded using 1-7 range) and that the value of the Slope feature should be in the 0-90 range. You can manually add these constraints to the auto-generated schema by setting the feature domain.
Load the auto-generated schema proto file
End of explanation
"""


schema_dir = os.path.join(ARTIFACT_STORE, "schema")
tf.io.gfile.makedirs(schema_dir)
schema_file = os.path.join(schema_dir, "schema.pbtxt")

tfdv.write_schema_text(schema, schema_file)

!cat {schema_file}

"""
Explanation: Modify the schema
You can use the protocol buffer APIs to modify the schema using tfdv.set_somain.
Review the TFDV library API documentation on setting a feature's domain. You can use the protocol buffer APIs to modify the schema. Review the Tensorflow Metadata proto definition for configuration options.
Save the updated schema
End of explanation
"""


schema_importer = ImporterNode(
    instance_name="Schema_Importer",
    source_uri=schema_dir,
    artifact_type=tfx.types.standard_artifacts.Schema,
    reimport=False,
)

context.run(schema_importer)

"""
Explanation: Importing the updated schema using ImporterNode
The ImporterNode component allows you to import an external artifact, including the schema file, so it can be used by other TFX components in your workflow. 
Configure and run the ImporterNode component
End of explanation
"""


context.show(schema_importer.outputs["result"])

"""
Explanation: Visualize the imported schema
End of explanation
"""


example_validator = ExampleValidator(
    instance_name="Data_Validation",
    statistics=statistics_gen.outputs["statistics"],
    schema=schema_importer.outputs["result"],
)

context.run(example_validator)

"""
Explanation: Validating data with ExampleValidator
The ExampleValidator component identifies anomalies in data.  It identifies anomalies by comparing data statistics computed by the StatisticsGen component against a schema generated by SchemaGen or imported by ImporterNode.
ExampleValidator can detect different classes of anomalies. For example it can:

perform validity checks by comparing data statistics against a schema 
detect training-serving skew by comparing training and serving data.
detect data drift by looking at a series of data.

The ExampleValidator component validates the data in the eval split only. Other splits are ignored. 
<img src=../../images/ExampleValidator.png width="350">
Configure and run the ExampleValidator component
End of explanation
"""


train_uri = example_validator.outputs["anomalies"].get()[0].uri
train_anomalies_filename = os.path.join(train_uri, "train/anomalies.pbtxt")
!cat $train_anomalies_filename

"""
Explanation: Examine the output of ExampleValidator
The output artifact of the ExampleValidator is the anomalies.pbtxt file describing an anomalies_pb2.Anomalies protobuf.
End of explanation
"""


context.show(example_validator.outputs["output"])

"""
Explanation: Visualize validation results
The file anomalies.pbtxt can be visualized using context.show.
End of explanation
"""


%%writefile config.py
FEATURE_KEY = "title"
LABEL_KEY = "source"
N_CLASSES = 3
HUB_URL = "https://tfhub.dev/google/nnlm-en-dim50/2"
HUB_DIM = 50
N_NEURONS = 16
TRAIN_BATCH_SIZE = 5
EVAL_BATCH_SIZE = 5
MODEL_NAME = "tfx_title_classifier"


def transformed_name(key):
    return key + "_xf"

%%writefile preprocessing.py
import tensorflow as tf
from config import FEATURE_KEY, LABEL_KEY, N_CLASSES, transformed_name


def _fill_in_missing(x):
    default_value = "" if x.dtype == tf.string else 0
    return tf.squeeze(
        tf.sparse.to_dense(
            tf.SparseTensor(x.indices, x.values, [x.dense_shape[0], 1]),
            default_value,
        ),
        axis=1,
    )


def preprocessing_fn(inputs):
    features = _fill_in_missing(inputs[FEATURE_KEY])
    labels = _fill_in_missing(inputs[LABEL_KEY])
    return {
        transformed_name(FEATURE_KEY): features,
        transformed_name(LABEL_KEY): labels,
    }

TRANSFORM_MODULE = "preprocessing.py"

"""
Explanation: In our case no anomalies were detected in the eval split.
For a detailed deep dive into data validation and schema generation refer to the lab-31-tfdv-structured-data lab.
Preprocessing data with Transform
The Transform component performs data transformation and feature engineering. The Transform component consumes tf.Examples emitted from the ExampleGen component and emits the transformed feature data and the SavedModel graph that was used to process the data. The emitted SavedModel  can then be used by serving components to make sure that the same data pre-processing logic is applied at training and serving.
The Transform component requires more code than many other components because of the arbitrary complexity of the feature engineering that you may need for the data and/or model that you're working with. It requires code files to be available which define the processing needed.
<img src=../../images/Transform.png width="400">
Define the pre-processing module
To configure Transform, you need to encapsulate your pre-processing code in the Python preprocessing_fn function and save it to a  python module that is then provided to the Transform component as an input. This module will be loaded by transform and the preprocessing_fn function will be called when the Transform component runs.
In most cases, your implementation of the preprocessing_fn makes extensive use of TensorFlow Transform for performing feature engineering on your dataset.
End of explanation
"""


transform = Transform(
    examples=example_gen.outputs["examples"],
    schema=schema_importer.outputs["result"],
    module_file=TRANSFORM_MODULE,
)

context.run(transform)

"""
Explanation: Configure and run the Transform component.
End of explanation
"""


os.listdir(transform.outputs["transform_graph"].get()[0].uri)

"""
Explanation: Examine the Transform component's outputs
The Transform component has 2 outputs:

transform_graph - contains the graph that can perform the preprocessing operations (this graph will be included in the serving and evaluation models).
transformed_examples - contains the preprocessed training and evaluation data.

Take a peek at the transform_graph artifact: it points to a directory containing 3 subdirectories:
End of explanation
"""


os.listdir(transform.outputs["transformed_examples"].get()[0].uri)

transform_uri = transform.outputs["transformed_examples"].get()[0].uri

tfrecord_filenames = [
    os.path.join(transform_uri, "train", name)
    for name in os.listdir(os.path.join(transform_uri, "train"))
]

dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type="GZIP")

for tfrecord in dataset.take(4):
    example = tf.train.Example()
    example.ParseFromString(tfrecord.numpy())
    for name, feature in example.features.feature.items():
        if feature.HasField("bytes_list"):
            value = feature.bytes_list.value
        if feature.HasField("float_list"):
            value = feature.float_list.value
        if feature.HasField("int64_list"):
            value = feature.int64_list.value
        print(f"{name}: {value}")
    print("******")

"""
Explanation: And the transform.examples artifact
End of explanation
"""


%%writefile model.py
import tensorflow as tf
import tensorflow_transform as tft
from config import (
    EVAL_BATCH_SIZE,
    HUB_DIM,
    HUB_URL,
    LABEL_KEY,
    MODEL_NAME,
    N_CLASSES,
    N_NEURONS,
    TRAIN_BATCH_SIZE,
    transformed_name,
)
from tensorflow.keras.callbacks import TensorBoard
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential
from tensorflow_hub import KerasLayer
from tfx_bsl.tfxio import dataset_options


def _get_serve_tf_examples_fn(model, tf_transform_output):
    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function
    def serve_tf_examples_fn(serialized_tf_examples):
        """Returns the output to be used in the serving signature."""
        feature_spec = tf_transform_output.raw_feature_spec()
        feature_spec.pop(LABEL_KEY)
        parsed_features = tf.io.parse_example(
            serialized_tf_examples, feature_spec
        )
        transformed_features = model.tft_layer(parsed_features)
        return model(transformed_features)

    return serve_tf_examples_fn


def _input_fn(file_pattern, data_accessor, tf_transform_output, batch_size=200):
    return data_accessor.tf_dataset_factory(
        file_pattern,
        dataset_options.TensorFlowDatasetOptions(
            batch_size=batch_size, label_key=transformed_name(LABEL_KEY)
        ),
        tf_transform_output.transformed_metadata.schema,
    )


def _load_hub_module_layer():
    hub_module = KerasLayer(
        HUB_URL,
        output_shape=[HUB_DIM],
        input_shape=[],
        dtype=tf.string,
        trainable=True,
    )
    return hub_module


def _build_keras_model():
    hub_module = _load_hub_module_layer()
    model = Sequential(
        [
            hub_module,
            Dense(N_NEURONS, activation="relu"),
            Dense(N_CLASSES, activation="softmax"),
        ]
    )
    model.compile(
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
    )
    return model


def run_fn(fn_args):

    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

    train_dataset = _input_fn(
        fn_args.train_files,
        fn_args.data_accessor,
        tf_transform_output,
        TRAIN_BATCH_SIZE,
    )

    eval_dataset = _input_fn(
        fn_args.eval_files,
        fn_args.data_accessor,
        tf_transform_output,
        EVAL_BATCH_SIZE,
    )

    mirrored_strategy = tf.distribute.MirroredStrategy()

    with mirrored_strategy.scope():
        model = _build_keras_model()

    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=fn_args.model_run_dir, update_freq="batch"
    )

    model.fit(
        train_dataset,
        steps_per_epoch=fn_args.train_steps,
        validation_data=eval_dataset,
        validation_steps=fn_args.eval_steps,
        callbacks=[tensorboard_callback],
    )

    signatures = {
        "serving_default": _get_serve_tf_examples_fn(
            model, tf_transform_output
        ).get_concrete_function(
            tf.TensorSpec(shape=[None], dtype=tf.string, name="examples")
        ),
    }
    model.save(
        fn_args.serving_model_dir, save_format="tf", signatures=signatures
    )

TRAINER_MODULE_FILE = "model.py"

"""
Explanation: Train your TensorFlow model with the Trainer component
The Trainer component trains a model using TensorFlow.
Trainer takes:

tf.Examples used for training and eval.
A user provided module file that defines the trainer logic.
A data schema created by SchemaGen or imported by ImporterNode.
A proto definition of train args and eval args.
An optional transform graph produced by upstream Transform component.
An optional base models used for scenarios such as warmstarting training.

<img src=../../images/Trainer.png width="400">
Define the trainer module
To configure Trainer, you need to encapsulate your training code in a Python module that is then provided to the Trainer as an input.
End of explanation
"""


trainer = Trainer(
    custom_executor_spec=executor_spec.ExecutorClassSpec(
        trainer_executor.GenericExecutor
    ),
    module_file=TRAINER_MODULE_FILE,
    transformed_examples=transform.outputs.transformed_examples,
    schema=schema_importer.outputs.result,
    transform_graph=transform.outputs.transform_graph,
    train_args=trainer_pb2.TrainArgs(splits=["train"], num_steps=20),
    eval_args=trainer_pb2.EvalArgs(splits=["eval"], num_steps=5),
)

context.run(trainer)

"""
Explanation: Create and run the Trainer component
As of the 0.25.0 release of TFX, the Trainer component only supports passing a single field - num_steps - through the train_args and eval_args arguments.
End of explanation
"""


logs_path = trainer.outputs["model_run"].get()[0].uri
print(logs_path)

"""
Explanation: Analyzing training runs with TensorBoard
In this step you will analyze the training run with TensorBoard.dev. TensorBoard.dev is a managed service that enables you to easily host, track and share your ML experiments.
Retrieve the location of TensorBoard logs
Each model run's train and eval metric logs are written to the model_run directory by the Tensorboard callback defined in model.py.
End of explanation
"""


model_resolver = ResolverNode(
    instance_name="latest_blessed_model_resolver",
    resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
    model=Channel(type=Model),
    model_blessing=Channel(type=ModelBlessing),
)

context.run(model_resolver)

"""
Explanation: Upload the logs and start TensorBoard.dev


Open a new JupyterLab terminal window


From the terminal window, execute the following command
tensorboard dev upload --logdir [YOUR_LOGDIR]


Where [YOUR_LOGDIR] is an URI retrieved by the previous cell.
You will be asked to authorize TensorBoard.dev using your Google account. If you don't have a Google account or you don't want to authorize TensorBoard.dev you can skip this exercise.
After the authorization process completes, follow the link provided to view your experiment.
Evaluating trained models with Evaluator
The Evaluator component analyzes model performance using the TensorFlow Model Analysis library. It runs inference requests on particular subsets of the test dataset, based on which slices are defined by the developer. Knowing which slices should be analyzed requires domain knowledge of what is important in this particular use case or domain. 
The Evaluator can also optionally validate a newly trained model against a previous model. In this lab, you only train one model, so the Evaluator automatically will label the model as "blessed".
<img src=../../images/Evaluator.png width="400">
Configure and run the Evaluator component
Use the ResolverNode to pick the previous model to compare against.  The model resolver is only required if performing model validation in addition to evaluation. In this case we validate against the latest blessed model. If no model has been blessed before (as in this case) the evaluator will make our candidate the first blessed model.
End of explanation
"""


accuracy_threshold = tfma.MetricThreshold(
    value_threshold=tfma.GenericValueThreshold(
        lower_bound={"value": 0.30}, upper_bound={"value": 0.99}
    )
)

metrics_specs = tfma.MetricsSpec(
    metrics=[
        tfma.MetricConfig(
            class_name="SparseCategoricalAccuracy", threshold=accuracy_threshold
        ),
        tfma.MetricConfig(class_name="ExampleCount"),
    ]
)

eval_config = tfma.EvalConfig(
    model_specs=[tfma.ModelSpec(label_key="source")],
    metrics_specs=[metrics_specs],
)
eval_config

model_analyzer = Evaluator(
    examples=example_gen.outputs.examples,
    model=trainer.outputs.model,
    baseline_model=model_resolver.outputs.model,
    eval_config=eval_config,
)

context.run(model_analyzer, enable_cache=False)

"""
Explanation: Configure evaluation metrics and slices.
End of explanation
"""


model_blessing_uri = model_analyzer.outputs.blessing.get()[0].uri
!ls -l {model_blessing_uri}

"""
Explanation: Check the model performance validation status
End of explanation
"""


trainer.outputs["model"]

pusher = Pusher(
    model=trainer.outputs["model"],
    model_blessing=model_analyzer.outputs["blessing"],
    push_destination=pusher_pb2.PushDestination(
        filesystem=pusher_pb2.PushDestination.Filesystem(
            base_directory=SERVING_MODEL_DIR
        )
    ),
)
context.run(pusher)

"""
Explanation: Deploying models with Pusher
The Pusher component checks whether a model has been "blessed", and if so, deploys it by pushing the model to a well known file destination.
<img src=../../images/Pusher.png width="400">
Configure and run the Pusher component
End of explanation
"""


pusher.outputs

# Set `PATH` to include a directory containing `saved_model_cli.
PATH = get_ipython().run_line_magic("env", "PATH")
%env PATH=/opt/conda/envs/tfx/bin:{PATH}

latest_pushed_model = os.path.join(
    SERVING_MODEL_DIR, max(os.listdir(SERVING_MODEL_DIR))
)
!saved_model_cli show --dir {latest_pushed_model} --all

"""
Explanation: Examine the output of Pusher
End of explanation
"""


