# Copyright 2021 The DDSP Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""Convert checkpoint to SavedModel, and SavedModel to TFJS / TFLite format.

Example Usage (Defaults):
ddsp_export --model_path=/path/to/model

Example Usage (TFJS model):
ddsp_export --model_path=/path/to/model --inference_model=autoencoder \
--tflite=false --tfjs

Example Usage (TFLite model):
ddsp_export --model_path=/path/to/model --inference_model=streaming_f0_pw \
--tflite --tfjs=false

Example Usage (SavedModel Only):
ddsp_export --model_path=/path/to/model --inference_model=[model_type] \
--tflite=false --tfjs=false
"""

import os

from absl import app
from absl import flags

from ddsp.training import inference
from ddsp.training import train_util
import gin
import tensorflow as tf
from tensorflowjs.converters import converter


from tflite_support import metadata as _metadata


flags.DEFINE_string(
    'model_path', '', 'Path to checkpoint or SavedModel directory. If no '
    'SavedModel is found, will search for latest checkpoint '
    'use it to create a SavedModel. Can also provide direct '
    'path to desired checkpoint. E.g. `/path/to/ckpt-[iter]`.')
flags.DEFINE_string(
    'save_dir', '', 'Optional directory in which to save converted checkpoint.'
    'If none is provided, it will be FLAGS.model_path if it '
    'contains a SavedModel, otherwise FLAGS.model_path/export.')

# Specify model class.
flags.DEFINE_enum(
    'inference_model',
    'streaming_f0_pw',
    [
        'autoencoder',
        'streaming_f0_pw',
    ],
    'Specify the ddsp.training.inference model to use for '
    'converting a checkpoint to a SavedModel. Names are '
    'snake_case versions of class names.')

# Optional flags.
flags.DEFINE_multi_string('gin_param', [],
                          'Gin parameters for custom inference model kwargs.')
flags.DEFINE_boolean('debug', False, 'DEBUG: Do not save the model')

# Conversion formats.
flags.DEFINE_boolean('tfjs', True,
                     'Convert SavedModel to TFJS for deploying on the web.')
flags.DEFINE_boolean('tflite', True,
                     'Convert SavedModel to TFLite for embedded C++ apps.')
flags.DEFINE_string('metadata_file', None,
                    'Optional metadata file to pack into TFLite model.')

FLAGS = flags.FLAGS


def get_inference_model(ckpt):
  """Restore model from checkpoint using global FLAGS.

  Use --gin_param for any custom kwargs for model constructors.
  Args:
    ckpt: Path to the checkpoint.

  Returns:
    Inference model, built and restored from checkpoint.
  """
  # Parse model kwargs from --gin_param.
  with gin.unlock_config():
    gin.parse_config_files_and_bindings(None, FLAGS.gin_param)

  models = {
      'autoencoder': inference.AutoencoderInference,
      'streaming_f0_pw': inference.StreamingF0PwInference,
  }
  return models[FLAGS.inference_model](ckpt)


def ckpt_to_saved_model(ckpt, save_dir):
  """Convert Checkpoint to SavedModel."""
  print(f'\nConverting to SavedModel:' f'\nInput: {ckpt}\nOutput: {save_dir}\n')
  model = get_inference_model(ckpt)
  print('Finshed Loading Model!')
  if not FLAGS.debug:
    model.save_model(save_dir)
  print('SavedModel Conversion Success!')


def saved_model_to_tfjs(input_dir, save_dir):
  """Convert SavedModel to TFJS model."""
  print(f'\nConverting to TFJS:\nInput:{input_dir}\nOutput:{save_dir}\n')
  converter.convert([
      '--input_format=tf_saved_model', '--signature_name=serving_default',
      '--control_flow_v2=True', '--skip_op_check', '--quantize_float16=True',
      '--experiments=True', input_dir, save_dir
  ])
  print('TFJS Conversion Success!')


def saved_model_to_tflite(input_dir, save_dir, metadata_file=None):
  """Convert SavedModel to TFLite model."""
  print(f'\nConverting to TFLite:\nInput:{input_dir}\nOutput:{save_dir}\n')
  # Convert the model.
  tflite_converter = tf.lite.TFLiteConverter.from_saved_model(input_dir)
  tflite_converter.target_spec.supported_ops = [
      tf.lite.OpsSet.TFLITE_BUILTINS,  # Enable TensorFlow Lite ops.
      tf.lite.OpsSet.SELECT_TF_OPS,  # Enable extended TensorFlow ops.
  ]
  tflite_model = tflite_converter.convert()  # Byte string.
  # Save the model.
  save_path = os.path.join(save_dir, 'model.tflite')
  with tf.io.gfile.GFile(save_path, 'wb') as f:
    f.write(tflite_model)

  if metadata_file is not None:
    populator = _metadata.MetadataPopulator.with_model_file(save_path)
    populator.load_associated_files([metadata_file])
    populator.populate()
  print('TFLite Conversion Success!')


def ensure_exits(dir_path):
  """Make directory if none exists."""
  if not tf.io.gfile.exists(dir_path):
    tf.io.gfile.makedirs(dir_path)


def main(unused_argv):
  model_path = train_util.expand_path(FLAGS.model_path)

  # Figure out what type the model path is.
  is_saved_model = tf.io.gfile.exists(
      os.path.join(model_path, 'saved_model.pb'))
  is_ckpt = not tf.io.gfile.isdir(model_path)

  # Infer save directory path.
  if FLAGS.save_dir:
    save_dir = FLAGS.save_dir
  else:
    if is_saved_model:
      # If model_path is a SavedModel, use that directory.
      save_dir = model_path
    elif is_ckpt:
      # If model_path is a checkpoint file, use the directory of the file.
      save_dir = os.path.join(os.path.dirname(model_path), 'export')
    else:
      # If model_path is a checkpoint directory, use child export directory.
      save_dir = os.path.join(model_path, 'export')

  # Make a new save directory.
  save_dir = train_util.expand_path(save_dir)
  ensure_exits(save_dir)

  # Create SavedModel if none already exists.
  if not is_saved_model:
    ckpt_to_saved_model(model_path, save_dir)

  # Convert SavedModel.
  if FLAGS.tfjs:
    tfjs_dir = os.path.join(save_dir, 'tfjs')
    ensure_exits(tfjs_dir)
    saved_model_to_tfjs(save_dir, tfjs_dir)

  if FLAGS.tflite:
    tflite_dir = os.path.join(save_dir, 'tflite')
    ensure_exits(tflite_dir)
    saved_model_to_tflite(save_dir, tflite_dir, FLAGS.metadata_file)


def console_entry_point():
  """From pip installed script."""
  app.run(main)


if __name__ == '__main__':
  console_entry_point()
