# 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.

# Lint as: python3
"""Plotting utilities for the DDSP library. Useful in colab and elsewhere."""

from ddsp import core
from ddsp import spectral_ops
from matplotlib import gridspec
import matplotlib.pyplot as plt
import numpy as np
import tensorflow.compat.v2 as tf

DEFAULT_SAMPLE_RATE = spectral_ops.CREPE_SAMPLE_RATE


def specplot(audio,
             vmin=-5,
             vmax=1,
             rotate=True,
             size=512 + 256,
             **matshow_kwargs):
  """Plot the log magnitude spectrogram of audio."""
  # If batched, take first element.
  if len(audio.shape) == 2:
    audio = audio[0]

  logmag = spectral_ops.compute_logmag(core.tf_float32(audio), size=size)
  if rotate:
    logmag = np.rot90(logmag)
  # Plotting.
  plt.matshow(logmag,
              vmin=vmin,
              vmax=vmax,
              cmap=plt.cm.magma,
              aspect='auto',
              **matshow_kwargs)
  plt.xticks([])
  plt.yticks([])
  plt.xlabel('Time')
  plt.ylabel('Frequency')


def transfer_function(ir, sample_rate=DEFAULT_SAMPLE_RATE):
  """Get true transfer function from an impulse_response."""
  n_fft = core.get_fft_size(0, ir.shape.as_list()[-1])
  frequencies = np.abs(
      np.fft.fftfreq(n_fft, 1 / sample_rate)[:int(n_fft / 2) + 1])
  magnitudes = tf.abs(tf.signal.rfft(ir, [n_fft]))
  return frequencies, magnitudes


def plot_impulse_responses(impulse_response,
                           desired_magnitudes,
                           sample_rate=DEFAULT_SAMPLE_RATE):
  """Plot a target frequency response, and that of an impulse response."""
  n_fft = desired_magnitudes.shape[-1] * 2
  frequencies = np.fft.fftfreq(n_fft, 1 / sample_rate)[:n_fft // 2]
  true_frequencies, true_magnitudes = transfer_function(impulse_response)

  # Plot it.
  plt.figure(figsize=(12, 6))
  plt.subplot(121)
  # Desired transfer function.
  plt.semilogy(frequencies, desired_magnitudes, label='Desired')
  # True transfer function.
  plt.semilogy(true_frequencies, true_magnitudes[0, 0, :], label='True')
  plt.title('Transfer Function')
  plt.legend()

  plt.subplot(122)
  plt.plot(impulse_response[0, 0, :])
  plt.title('Impulse Response')


def pianoroll_plot_setup(figsize=None, side_piano_ratio=0.025,
                         faint_pr=True, xlim=None):
  """Makes a tiny piano left of the y-axis and a faint piano on the main figure.

  This function sets up the figure for pretty plotting a piano roll. It makes a
  small imshow plot to the left of the main plot that looks like a piano. This
  piano side plot is aligned along the y-axis of the main plot, such that y
  values align with MIDI values (y=0 is the lowest C-1, y=11 is C0, etc).
  Additionally, a main figure is set up that shares the y-axis of the piano side
  plot. Optionally, a set of faint horizontal lines are drawn on the main figure
  that correspond to the black keys on the piano (and a line separating B & C
  and E & F). This function returns the formatted figure, the side piano axis,
  and the main axis for plotting your data.

  By default, this will draw 11 octaves of piano keys along the y-axis; you will
  probably want reduce what is visible using `ax.set_ylim()` on either returned
  axis.

  Using with imshow piano roll data:
    A common use case is for using imshow() on the main axis to display a piano
    roll alongside the piano side plot AND the faint piano roll behind your
    data. In this case, if your data is a 2D array you have to use a masked
    numpy array to make certain values invisible on the plot, and therefore make
    the faint piano roll visible. Here's an example:

    midi = np.flipud([
          [0.0, 0.0, 1.0],
          [0.0, 1.0, 0.0],
          [1.0, 0.0, 0.0],
    ])

    midi_masked = np.ma.masked_values(midi, 0.0)  # Mask out all 0.0's
    fig, ax, sp = plotting.pianoroll_plot_setup()
    ax.imshow(midi_masked, origin='lower', aspect='auto')  # main subplot axis
    sp.set_ylabel('My favorite MIDI data')  # side piano axis
    fig.show()

    The other option is to use imshow in RGBA mode, where your data is split
    into 4 channels. Every alpha value that is 0.0 will be transparent and show
    the faint piano roll below your data.

  Args:
    figsize: Size if the matplotlib figure. Will be passed to `plt.figure()`.
      Defaults to None.
    side_piano_ratio: Width of the y-axis piano in terms of raio of the whole
      figure. Defaults to 1/40th.
    faint_pr: Whether to draw faint black & white keys across the main plot.
      Defaults to True.
    xlim: Tuple containing the min and max of the x values for the main plot.
      Only used to determine the x limits for the faint piano roll in the main
      plot. Defaults to (0, 1000).

  Returns:
    (figure, main_axis, left_piano_axis)
      figure: A matplotlib figure object containing both subplots set up with an
        aligned piano roll.
      main_axis: A matplotlib axis object to be used for plotting. Optionally
        has a faint piano roll in the background.
      left_piano_axis: A matplotlib axis object that has a small, aligned piano
        along the left side y-axis of the main_axis subplot.
  """
  octaves = 11

  # Setup figure and gridspec.
  fig = plt.figure(figsize=figsize)
  gs_ratio = int(1 / side_piano_ratio)
  gs = gridspec.GridSpec(1, 2, width_ratios=[1, gs_ratio])
  left_piano_ax = fig.add_subplot(gs[0])

  # Make a piano on the left side of the y-axis with imshow().
  keys = np.array(
      [0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0]  # notes in descending order; B -> C
  )
  keys = np.tile(keys, octaves)[:, None]
  left_piano_ax.imshow(keys, cmap='binary', aspect='auto',
                       extent=[0, 0.625, -0.5, octaves*12-0.5])

  # Make the lines between keys.
  for i in range(octaves):
    left_piano_ax.hlines(i*12 - 0.5, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 1.0, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 3.0, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 4.5, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 6.0, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 8.0, -0.5, 1, colors='black', linewidth=0.5)
    left_piano_ax.hlines(i*12 + 10.0, -0.5, 1, colors='black', linewidth=0.5)

  # Set the limits of the side piano and remove ticks so it looks nice.
  left_piano_ax.set_xlim(0, 0.995)
  left_piano_ax.set_xticks([])

  # Create the aligned axis we'll return to the user.
  main_ax = fig.add_subplot(gs[1], sharey=left_piano_ax)

  # Draw a faint piano roll behind the main axes (if the user wants).
  if faint_pr:
    xlim = (0, 1000) if xlim is None else xlim
    x_min, x_max = xlim
    x_delta = x_max - x_min
    main_ax.imshow(np.tile(keys, x_delta), cmap='binary', aspect='auto',
                   alpha=0.05, extent=[x_min, x_max, -0.5, octaves*12-0.5])
    for i in range(octaves):
      main_ax.hlines(i * 12 + 4.5, x_min, x_max, colors='black',
                     linewidth=0.5, alpha=0.25)
      main_ax.hlines(i * 12 - 0.5, x_min, x_max, colors='black',
                     linewidth=0.5, alpha=0.25)

    main_ax.set_xlim(*xlim)

  # Some final cosmetic tweaks before returning the axis obj's and figure.
  plt.setp(main_ax.get_yticklabels(), visible=False)
  gs.tight_layout(fig)
  return fig, main_ax, left_piano_ax
