# coding=utf-8
# Copyright 2022 The Tensor2Tensor 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.

"""Hyperparameters defining different problems.

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.layers import modalities
from tensor2tensor.utils import registry

import tensorflow.compat.v1 as tf

# TODO(rsepassi): Merge these problems with their data generators. Currently
# they only implement the hparams.


class AudioTimitProblem(problem.Problem):
  """Base class for TIMIT problems."""

  def example_reading_spec(self):
    data_fields = {
        "inputs": tf.VarLenFeature(tf.int64),
        "audio/sample_count": tf.FixedLenFeature((), tf.int64),
        "audio/sample_width": tf.FixedLenFeature((), tf.int64),
        "targets": tf.VarLenFeature(tf.int64),
    }
    return data_fields, None

  def preprocess_example(self, example, mode, hparams):
    example = super(AudioTimitProblem, self).preprocess_example(
        example, mode, hparams)
    # Reshape audio to proper shape
    sample_count = tf.to_int32(example.pop("audio/sample_count"))
    sample_width = tf.to_int32(example.pop("audio/sample_width"))
    channel_count = 1
    example["inputs"] = tf.reshape(example["inputs"],
                                   [sample_count, sample_width, channel_count])
    return example


@registry.register_problem
class AudioTimitCharactersTune(AudioTimitProblem):
  """TIMIT to characters."""

  def feature_encoders(self, _):
    return {
        "inputs": text_encoder.TextEncoder(),
        "targets": text_encoder.ByteTextEncoder(),
    }

  def hparams(self, defaults, model_hparams):
    hp = defaults
    hp.modality = {"inputs": modalities.ModalityType.SPEECH_RECOGNITION,
                   "targets": modalities.ModalityType.SYMBOL}
    hp.vocab_size = {"inputs": None,
                     "targets": 256}


@registry.register_problem
class AudioTimitTokens8kTune(AudioTimitProblem):
  """TIMIT to tokens."""

  @property
  def target_vocab_size(self):
    return 2**13  # 8192

  def feature_encoders(self, data_dir):
    vocab_filename = os.path.join(data_dir,
                                  "vocab.endefr.%d" % self.target_vocab_size)
    subtokenizer = text_encoder.SubwordTextEncoder(vocab_filename)
    return {
        "inputs": text_encoder.TextEncoder(),
        "targets": subtokenizer,
    }