import gin
import tensorflow as tf
import tensorflow.keras.layers as tfkl
from ddsp.training.models import Autoencoder


@gin.register
class VectorQuantizer(tfkl.Layer):
    def __init__(self, num_embeddings, initializer='uniform', **kwargs):
        self.num_embeddings = num_embeddings
        self.initializer = initializer
        self.embedding_dim = None
        self.codebook = None
        super(VectorQuantizer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Add embedding weights.
        self.embedding_dim = input_shape['z'][-1]
        self.codebook = self.add_weight(name='codebook',
                                        shape=(self.num_embeddings, self.embedding_dim),
                                        initializer=self.initializer,
                                        trainable=True)
        # Finalize building.
        super(VectorQuantizer, self).build(input_shape)

    def call(self, conditioning, **kwargs):
        encodings = conditioning['z']

        encodings_flat = tf.reshape(encodings, shape=(-1, self.embedding_dim))
        distances = tf.reduce_mean((encodings_flat[:, None] - self.codebook[None]) ** 2, axis=-1)
        min_indices = tf.argmin(distances, axis=-1)
        codes = tf.gather(self.codebook, min_indices, axis=0)
        codes = tf.reshape(codes, shape=[-1, *encodings.shape[1:]])

        # Straight through estimator trick
        # WARNING: codes_forward will not be equal to codes_for_loss, presumably due to numerical instabilities
        codes_forward, codes_for_loss = encodings + tf.stop_gradient(codes - encodings), codes
        conditioning['z_preq'] = conditioning['z']
        conditioning['z'] = codes_forward
        conditioning['z_for_loss'] = codes_for_loss
        conditioning['z_indices'] = min_indices
#         print(conditioning.keys())
        return conditioning


@gin.configurable
class QuantizingAutoencoder(Autoencoder):
    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'], conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):
    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    # noinspection PyMethodOverriding,PyUnresolvedReferences
    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    # noinspection PyMethodOverriding
    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    # noinspection PyMethodOverriding
    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)

a2550b783bf76cc745cb5cad583400f45163ffcd
<mask token>


@gin.register
class VectorQuantizer(tfkl.Layer):
    <mask token>
    <mask token>
    <mask token>


@gin.configurable
class QuantizingAutoencoder(Autoencoder):

    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'],
                        conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):

    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)

<mask token>


@gin.register
class VectorQuantizer(tfkl.Layer):
    <mask token>

    def build(self, input_shape):
        self.embedding_dim = input_shape['z'][-1]
        self.codebook = self.add_weight(name='codebook', shape=(self.
            num_embeddings, self.embedding_dim), initializer=self.
            initializer, trainable=True)
        super(VectorQuantizer, self).build(input_shape)
    <mask token>


@gin.configurable
class QuantizingAutoencoder(Autoencoder):

    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'],
                        conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):

    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)

<mask token>


@gin.register
class VectorQuantizer(tfkl.Layer):

    def __init__(self, num_embeddings, initializer='uniform', **kwargs):
        self.num_embeddings = num_embeddings
        self.initializer = initializer
        self.embedding_dim = None
        self.codebook = None
        super(VectorQuantizer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.embedding_dim = input_shape['z'][-1]
        self.codebook = self.add_weight(name='codebook', shape=(self.
            num_embeddings, self.embedding_dim), initializer=self.
            initializer, trainable=True)
        super(VectorQuantizer, self).build(input_shape)
    <mask token>


@gin.configurable
class QuantizingAutoencoder(Autoencoder):

    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'],
                        conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):

    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)

import gin
import tensorflow as tf
import tensorflow.keras.layers as tfkl
from ddsp.training.models import Autoencoder


@gin.register
class VectorQuantizer(tfkl.Layer):

    def __init__(self, num_embeddings, initializer='uniform', **kwargs):
        self.num_embeddings = num_embeddings
        self.initializer = initializer
        self.embedding_dim = None
        self.codebook = None
        super(VectorQuantizer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.embedding_dim = input_shape['z'][-1]
        self.codebook = self.add_weight(name='codebook', shape=(self.
            num_embeddings, self.embedding_dim), initializer=self.
            initializer, trainable=True)
        super(VectorQuantizer, self).build(input_shape)

    def call(self, conditioning, **kwargs):
        encodings = conditioning['z']
        encodings_flat = tf.reshape(encodings, shape=(-1, self.embedding_dim))
        distances = tf.reduce_mean((encodings_flat[:, None] - self.codebook
            [None]) ** 2, axis=-1)
        min_indices = tf.argmin(distances, axis=-1)
        codes = tf.gather(self.codebook, min_indices, axis=0)
        codes = tf.reshape(codes, shape=[-1, *encodings.shape[1:]])
        codes_forward, codes_for_loss = encodings + tf.stop_gradient(codes -
            encodings), codes
        conditioning['z_preq'] = conditioning['z']
        conditioning['z'] = codes_forward
        conditioning['z_for_loss'] = codes_for_loss
        conditioning['z_indices'] = min_indices
        return conditioning


@gin.configurable
class QuantizingAutoencoder(Autoencoder):

    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'],
                        conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):

    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)

import gin
import tensorflow as tf
import tensorflow.keras.layers as tfkl
from ddsp.training.models import Autoencoder


@gin.register
class VectorQuantizer(tfkl.Layer):
    def __init__(self, num_embeddings, initializer='uniform', **kwargs):
        self.num_embeddings = num_embeddings
        self.initializer = initializer
        self.embedding_dim = None
        self.codebook = None
        super(VectorQuantizer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Add embedding weights.
        self.embedding_dim = input_shape['z'][-1]
        self.codebook = self.add_weight(name='codebook',
                                        shape=(self.num_embeddings, self.embedding_dim),
                                        initializer=self.initializer,
                                        trainable=True)
        # Finalize building.
        super(VectorQuantizer, self).build(input_shape)

    def call(self, conditioning, **kwargs):
        encodings = conditioning['z']

        encodings_flat = tf.reshape(encodings, shape=(-1, self.embedding_dim))
        distances = tf.reduce_mean((encodings_flat[:, None] - self.codebook[None]) ** 2, axis=-1)
        min_indices = tf.argmin(distances, axis=-1)
        codes = tf.gather(self.codebook, min_indices, axis=0)
        codes = tf.reshape(codes, shape=[-1, *encodings.shape[1:]])

        # Straight through estimator trick
        # WARNING: codes_forward will not be equal to codes_for_loss, presumably due to numerical instabilities
        codes_forward, codes_for_loss = encodings + tf.stop_gradient(codes - encodings), codes
        conditioning['z_preq'] = conditioning['z']
        conditioning['z'] = codes_forward
        conditioning['z_for_loss'] = codes_for_loss
        conditioning['z_indices'] = min_indices
#         print(conditioning.keys())
        return conditioning


@gin.configurable
class QuantizingAutoencoder(Autoencoder):
    def __init__(self, num_embeddings=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.quantizer = VectorQuantizer(num_embeddings)

    def call(self, features, training=True):
        """Run the core of the network, get predictions and loss."""
        conditioning = self.encode(features, training=training)
        conditioning = self.quantizer(conditioning)
        audio_gen = self.decode(conditioning, training=training)
        if training:
            for loss_obj in self.loss_objs:
                if isinstance(loss_obj, QuantizationLoss):
                    loss = loss_obj(conditioning['z_for_loss'], conditioning['z_preq'])
                else:
                    loss = loss_obj(features['audio'], audio_gen)
                self._losses_dict[loss_obj.name] = loss
        return audio_gen


@gin.register
class QuantizationLoss(tfkl.Layer):
    def __init__(self, weight, name):
        super().__init__(name=name)
        self.weight = weight

    # noinspection PyMethodOverriding,PyUnresolvedReferences
    def call(self, codes, encodings):
        return self.calc_loss(codes, encodings) * self.weight

    def calc_loss(self, codes, encodings):
        pass


@gin.register
class CodebookLoss(QuantizationLoss):

    def __init__(self, weight, name='codebook_loss'):
        super().__init__(weight=weight, name=name)

    # noinspection PyMethodOverriding
    def calc_loss(self, codes, encodings):
        """
        Calculates the codebook loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((tf.stop_gradient(encodings) - codes) ** 2)


@gin.register
class CommitmentLoss(QuantizationLoss):

    def __init__(self, weight, name='commitment_loss'):
        super().__init__(weight=weight, name=name)

    # noinspection PyMethodOverriding
    def calc_loss(self, codes, encodings):
        """
        Calculates the commitment loss.
        :param codes: The discrete codes that are the result of the quantization process.
        :param encodings: The original non-discrete values that were generated by the encoder.
        :return: The loss.
        """
        return tf.reduce_mean((encodings - tf.stop_gradient(codes)) ** 2)
