# AUTOGENERATED! DO NOT EDIT! File to edit: tensorflow-tflite_converter.ipynb (unless otherwise specified).

__all__ = ['TFLITE_MODEL_DEFAULT_NAME', 'logger', 'convert_model', 'save_model', 'read_pipeline_config',
           'configure_logging']


# Cell
import logging
import logging.handlers
import argparse
import sys
import tensorflow as tf
from os.path import join
from google.protobuf import text_format
from ..dataset.type import DatasetType, infer_dataset_type
from .tflite_metadata import create_metadata, write_metadata


# Cell
TFLITE_MODEL_DEFAULT_NAME = 'model.tflite'


# Cell
logger = logging.getLogger(__name__)


# Cell
def convert_model(saved_model_dir):
    """
    Convert a SavedModel into Tensorflow Lite Format.
    `saved_model_dir`: the path to the SavedModel directory
    returns: the converted Tensorflow Lite model
    """
    logger.info('Converting SavedModel from: {}'.format(saved_model_dir))
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)  # path to the SavedModel directory
    tflite_model = converter.convert()
    return tflite_model


# Cell
def save_model(tflite_model, output_file):
    """
    Save a Tensowflow Lite model to disk.
    `tflite_model`: the Tensorflow Lite model
    `output_file`: the path and filename to save the Tensorflow Lite model
    """
    with open(output_file, 'wb') as f:
        f.write(tflite_model)
    logger.info('Successfully save model to file: {}'.format(output_file))


# Cell
def read_pipeline_config(pipeline_config_path):
    """
    Reads the pipeline config file.

    `pipeline_config_path`: The path to the pipeline config file.
    """
    pipeline_config = {}
    with tf.io.gfile.GFile(pipeline_config_path, 'r') as f:
        text_format.Parse(f.read(), pipeline_config)
    return pipeline_config


# Cell
def configure_logging(logging_level=logging.INFO):
    """
    Configures logging for the system.

    `logging_level`: The logging level to use.
    """
    logger.setLevel(logging_level)

    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(logging_level)

    logger.addHandler(handler)


# Cell
if __name__ == '__main__' and '__file__' in globals():
    configure_logging()

    parser = argparse.ArgumentParser()
    parser.add_argument("-s",
                        "--source",
                        help="The path to the folder containing the SavedModel.",
                        type=str)
    parser.add_argument("-c",
                        "--categories",
                        help="The categories file to add to the Tensorflow Lite model.",
                        type=str)
    parser.add_argument("-n",
                        "--name",
                        help="The name of the model.",
                        type=str)
    parser.add_argument("-v",
                        "--version",
                        help="The version of the model.",
                        type=int,
                        default=1)
    parser.add_argument("-t",
                        "--type",
                        help="The type of the model, if not explicitly set try to infer from categories file path.",
                        choices=list(DatasetType),
                        type=DatasetType,
                        default=None)
    parser.add_argument("-o",
                        "--output",
                        help="The folder to store the Tensorflow Lite model.",
                        type=str)
    args = parser.parse_args()

    model_type = args.type

    # try to infer the model type if not explicitly set
    if model_type is None:
        try:
            model_type = infer_dataset_type(args.categories)
        except ValueError as e:
            logger.error(e)
            sys.exit(1)

    output_file = join(args.output, TFLITE_MODEL_DEFAULT_NAME)

    save_model(convert_model(args.source), output_file)

    model_meta = create_metadata(args.source, args.categories, model_type, args.name, args.version)
    write_metadata(model_meta, output_file, args.categories)

    logger.info('FINISHED!!!')
