# Convert to mobile formats
import coremltools
import tensorflow as tf
import tempfile

def convert_to_coreml(model):
    return coremltools.converters.keras.convert(
        model,
        input_names=['input'],
        output_names=['digit']
    )

def convert_to_tflite(model):
    # save the model to a temp file so we can
    # convert it.
    keras_file = tempfile.mktemp()
    model.save(keras_file, include_optimizer=False)
    converter = tf.lite.TFLiteConverter.from_keras_model_file(keras_file)
    return converter.convert()

mlmodel = convert_to_coreml(model)
tflite = convert_to_tflite(model)

model_name = "mnist_cnn_lr001_batchsize128"

# Keras
model.save(
    model_name + ".h5",
    include_optimizer="False"
)

# Core ML
mlmodel.save(model_name + ".mlmodel")


# TensorFlow Lite
with open(model_name + ".tflite", 'wb') as wfid:
    wfid.write(tflite)