# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================

"""Sample data exhibiting audio summaries, via a waveform generator."""

import inspect
import math
import os.path

from absl import app
from absl import flags
import tensorflow as tf
from tensorboard.plugins.audio import summary

tf.compat.v1.disable_eager_execution()

FLAGS = flags.FLAGS

flags.DEFINE_string(
    "logdir",
    "/tmp/audio_demo",
    "Directory into which to write TensorBoard data.",
)

flags.DEFINE_integer("steps", 50, "Number of frequencies of each waveform to generate.")

# Parameters for the audio output.
flags.DEFINE_integer("sample_rate", 44100, "Sample rate, in Hz.")
flags.DEFINE_float("duration", 2.0, "Duration of each waveform, in s.")


def _samples():
    """Compute how many samples should be included in each waveform."""
    return int(FLAGS.sample_rate * FLAGS.duration)


def run(logdir, run_name, wave_name, wave_constructor):
    """Generate wave data of the given form.

    The provided function `wave_constructor` should accept a scalar tensor
    of type float32, representing the frequency (in Hz) at which to
    construct a wave, and return a tensor of shape [1, _samples(), `n`]
    representing audio data (for some number of channels `n`).

    Waves will be generated at frequencies ranging from A4 to A5.

    Arguments:
      logdir: the top-level directory into which to write summary data
      run_name: the name of this run; will be created as a subdirectory
        under logdir
      wave_name: the name of the wave being generated
      wave_constructor: see above
    """
    tf.compat.v1.reset_default_graph()
    tf.compat.v1.set_random_seed(0)

    # On each step `i`, we'll set this placeholder to `i`. This allows us
    # to know "what time it is" at each step.
    step_placeholder = tf.compat.v1.placeholder(tf.float32, shape=[])

    # We want to linearly interpolate a frequency between A4 (440 Hz) and
    # A5 (880 Hz).
    with tf.name_scope("compute_frequency"):
        f_min = 440.0
        f_max = 880.0
        t = step_placeholder / (FLAGS.steps - 1)
        frequency = f_min * (1.0 - t) + f_max * t

    # Let's log this frequency, just so that we can make sure that it's as
    # expected.
    tf.compat.v1.summary.scalar("frequency", frequency)

    # Now, we pass this to the wave constructor to get our waveform. Doing
    # so within a name scope means that any summaries that the wave
    # constructor produces will be namespaced.
    with tf.name_scope(wave_name):
        waveform = wave_constructor(frequency)

    # We also have the opportunity to annotate each audio clip with a
    # label. This is a good place to include the frequency, because it'll
    # be visible immediately next to the audio clip.
    with tf.name_scope("compute_labels"):
        samples = tf.shape(input=waveform)[0]
        wave_types = tf.tile(["*Wave type:* `%s`." % wave_name], [samples])
        frequencies = tf.strings.join(
            [
                "*Frequency:* ",
                tf.tile([tf.as_string(frequency, precision=2)], [samples]),
                " Hz.",
            ]
        )
        samples = tf.strings.join(
            [
                "*Sample:* ",
                tf.as_string(tf.range(samples) + 1),
                " of ",
                tf.as_string(samples),
                ".",
            ]
        )
        labels = tf.strings.join([wave_types, frequencies, samples], separator=" ")

    # We can place a description next to the summary in TensorBoard. This
    # is a good place to explain what the summary represents, methodology
    # for creating it, etc. Let's include the source code of the function
    # that generated the wave.
    source = "\n".join(
        "    %s" % line.rstrip() for line in inspect.getsourcelines(wave_constructor)[0]
    )
    description = "A wave of type `%r`, generated via:\n\n%s" % (
        wave_name,
        source,
    )

    # Here's the crucial piece: we interpret this result as audio.
    summary.op(
        "waveform",
        waveform,
        FLAGS.sample_rate,
        labels=labels,
        display_name=wave_name,
        description=description,
    )

    # Now, we can collect up all the summaries and begin the run.
    summ = tf.compat.v1.summary.merge_all()

    sess = tf.compat.v1.Session()
    writer = tf.compat.v1.summary.FileWriter(os.path.join(logdir, run_name))
    writer.add_graph(sess.graph)
    sess.run(tf.compat.v1.global_variables_initializer())
    for step in range(FLAGS.steps):
        s = sess.run(summ, feed_dict={step_placeholder: float(step)})
        writer.add_summary(s, global_step=step)
    writer.close()


# Now, let's take a look at the kinds of waves that we can generate.


def sine_wave(frequency):
    """Emit a sine wave at the given frequency."""
    xs = tf.reshape(tf.range(_samples(), dtype=tf.float32), [1, _samples(), 1])
    ts = xs / FLAGS.sample_rate
    return tf.sin(2 * math.pi * frequency * ts)


def square_wave(frequency):
    """Emit a square wave at the given frequency."""
    # The square is just the sign of the sine!
    return tf.sign(sine_wave(frequency))


def triangle_wave(frequency):
    """Emit a triangle wave at the given frequency."""
    xs = tf.reshape(tf.range(_samples(), dtype=tf.float32), [1, _samples(), 1])
    ts = xs / FLAGS.sample_rate
    #
    # A triangle wave looks like this:
    #
    #      /\      /\
    #     /  \    /  \
    #         \  /    \  /
    #          \/      \/
    #
    # If we look at just half a period (the first four slashes in the
    # diagram above), we can see that it looks like a transformed absolute
    # value function.
    #
    # Let's start by computing the times relative to the start of each
    # half-wave pulse (each individual "mountain" or "valley", of which
    # there are four in the above diagram).
    half_pulse_index = ts * (frequency * 2)
    half_pulse_angle = half_pulse_index % 1.0  # in [0, 1]
    #
    # Now, we can see that each positive half-pulse ("mountain") has
    # amplitude given by A(z) = 0.5 - abs(z - 0.5), and then normalized:
    absolute_amplitude = (0.5 - tf.abs(half_pulse_angle - 0.5)) / 0.5
    #
    # But every other half-pulse is negative, so we should invert these.
    half_pulse_parity = tf.sign(1 - (half_pulse_index % 2.0))
    amplitude = half_pulse_parity * absolute_amplitude
    #
    # This is precisely the desired result, so we're done!
    return amplitude


# If we want to get fancy, we can use our above waves as primitives to
# build more interesting waves.


def bisine_wave(frequency):
    """Emit two sine waves, in stereo at different octaves."""
    #
    # We can first our existing sine generator to generate two different
    # waves.
    f_hi = frequency
    f_lo = frequency / 2.0
    with tf.name_scope("hi"):
        sine_hi = sine_wave(f_hi)
    with tf.name_scope("lo"):
        sine_lo = sine_wave(f_lo)
    #
    # Now, we have two tensors of shape [1, _samples(), 1]. By concatenating
    # them along axis 2, we get a tensor of shape [1, _samples(), 2]---a
    # stereo waveform.
    return tf.concat([sine_lo, sine_hi], axis=2)


def bisine_wahwah_wave(frequency):
    """Emit two sine waves with balance oscillating left and right."""
    #
    # This is clearly intended to build on the bisine wave defined above,
    # so we can start by generating that.
    waves_a = bisine_wave(frequency)
    #
    # Then, by reversing axis 2, we swap the stereo channels. By mixing
    # this with `waves_a`, we'll be able to create the desired effect.
    waves_b = tf.reverse(waves_a, axis=[2])
    #
    # Let's have the balance oscillate from left to right four times.
    iterations = 4
    #
    # Now, we compute the balance for each sample: `ts` has values
    # in [0, 1] that indicate how much we should use `waves_a`.
    xs = tf.reshape(tf.range(_samples(), dtype=tf.float32), [1, _samples(), 1])
    thetas = xs / _samples() * iterations
    ts = (tf.sin(math.pi * 2 * thetas) + 1) / 2
    #
    # Finally, we can mix the two together, and we're done.
    wave = ts * waves_a + (1.0 - ts) * waves_b
    #
    # Alternately, we can make the effect more pronounced by exaggerating
    # the sample data. Let's emit both variations.
    exaggerated_wave = wave**3.0
    return tf.concat([wave, exaggerated_wave], axis=0)


def run_all(logdir, verbose=False):
    """Generate waves of the shapes defined above.

    Arguments:
      logdir: the directory into which to store all the runs' data
      verbose: if true, print out each run's name as it begins
    """
    waves = [
        sine_wave,
        square_wave,
        triangle_wave,
        bisine_wave,
        bisine_wahwah_wave,
    ]
    for i, wave_constructor in enumerate(waves):
        wave_name = wave_constructor.__name__
        run_name = "%02d_%s" % (i + 1, wave_name)
        if verbose:
            print("--- Running: %s" % run_name)
        run(logdir, run_name, wave_name, wave_constructor)


def main(unused_argv):
    print("Saving output to %s." % FLAGS.logdir)
    run_all(FLAGS.logdir, verbose=True)
    print("Done. Output saved to %s." % FLAGS.logdir)


if __name__ == "__main__":
    app.run(main)
