# Copyright 2020 The Magenta 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.

"""Performance RNN model."""

import collections
import functools

from magenta.contrib import training as contrib_training
from magenta.models.shared import events_rnn_model
import note_seq
from note_seq.protobuf import generator_pb2

# State for constructing a time-varying control sequence. Keeps track of the
# current event position and time step in the generated performance, to allow
# the control sequence to vary with clock time.
PerformanceControlState = collections.namedtuple(
    'PerformanceControlState', ['current_perf_index', 'current_perf_step'])


class PerformanceRnnModel(events_rnn_model.EventSequenceRnnModel):
  """Class for RNN performance generation models."""

  def generate_performance(
      self, num_steps, primer_sequence, temperature=1.0, beam_size=1,
      branch_factor=1, steps_per_iteration=1, control_signal_fns=None,
      disable_conditioning_fn=None):
    """Generate a performance track from a primer performance track.

    Args:
      num_steps: The integer length in steps of the final track, after
          generation. Includes the primer.
      primer_sequence: The primer sequence, a Performance object.
      temperature: A float specifying how much to divide the logits by
         before computing the softmax. Greater than 1.0 makes tracks more
         random, less than 1.0 makes tracks less random.
      beam_size: An integer, beam size to use when generating tracks via
          beam search.
      branch_factor: An integer, beam search branch factor to use.
      steps_per_iteration: An integer, number of steps to take per beam search
          iteration.
      control_signal_fns: A list of functions that map time step to desired
          control value, or None if not using control signals.
      disable_conditioning_fn: A function that maps time step to whether or not
          conditioning should be disabled, or None if there is no conditioning
          or conditioning is not optional.

    Returns:
      The generated Performance object (which begins with the provided primer
      track).
    """
    if control_signal_fns:
      control_event = tuple(f(0) for f in control_signal_fns)
      if disable_conditioning_fn is not None:
        control_event = (disable_conditioning_fn(0), control_event)
      control_events = [control_event]
      control_state = PerformanceControlState(
          current_perf_index=0, current_perf_step=0)
      extend_control_events_callback = functools.partial(
          _extend_control_events, control_signal_fns, disable_conditioning_fn)
    else:
      control_events = None
      control_state = None
      extend_control_events_callback = None

    return self._generate_events(
        num_steps, primer_sequence, temperature, beam_size, branch_factor,
        steps_per_iteration, control_events=control_events,
        control_state=control_state,
        extend_control_events_callback=extend_control_events_callback)

  def performance_log_likelihood(self, sequence, control_values,
                                 disable_conditioning):
    """Evaluate the log likelihood of a performance.

    Args:
      sequence: The Performance object for which to evaluate the log likelihood.
      control_values: List of (single) values for all control signal.
      disable_conditioning: Whether or not to disable optional conditioning. If
          True, disable conditioning. If False, do not disable. None when no
          conditioning or it is not optional.


    Returns:
      The log likelihood of `sequence` under this model.
    """
    if control_values:
      control_event = tuple(control_values)
      if disable_conditioning is not None:
        control_event = (disable_conditioning, control_event)
      control_events = [control_event] * len(sequence)
    else:
      control_events = None

    return self._evaluate_log_likelihood(
        [sequence], control_events=control_events)[0]


def _extend_control_events(control_signal_fns, disable_conditioning_fn,
                           control_events, performance, control_state):
  """Extend a performance control sequence.

  Extends `control_events` -- a sequence of control signal value tuples -- to be
  one event longer than `performance`, so the next event of `performance` can be
  conditionally generated.

  This function is meant to be used as the `extend_control_events_callback`
  in the `_generate_events` method of `EventSequenceRnnModel`.

  Args:
    control_signal_fns: A list of functions that map time step to desired
        control value, or None if not using control signals.
    disable_conditioning_fn: A function that maps time step to whether or not
        conditioning should be disabled, or None if there is no conditioning or
        conditioning is not optional.
    control_events: The control sequence to extend.
    performance: The Performance being generated.
    control_state: A PerformanceControlState tuple containing the current
        position in `performance`. We maintain this so as not to have to
        recompute the total performance length (in steps) every time we want to
        extend the control sequence.

  Returns:
    The PerformanceControlState after extending the control sequence one step
    past the end of the generated performance.
  """
  idx = control_state.current_perf_index
  step = control_state.current_perf_step

  while idx < len(performance):
    if performance[idx].event_type == note_seq.PerformanceEvent.TIME_SHIFT:
      step += performance[idx].event_value
    idx += 1

    control_event = tuple(f(step) for f in control_signal_fns)
    if disable_conditioning_fn is not None:
      control_event = (disable_conditioning_fn(step), control_event)
    control_events.append(control_event)

  return PerformanceControlState(
      current_perf_index=idx, current_perf_step=step)


class PerformanceRnnConfig(events_rnn_model.EventSequenceRnnConfig):
  """Stores a configuration for a Performance RNN.

  Attributes:
    num_velocity_bins: Number of velocity bins to use. If 0, don't use velocity
        at all.
    control_signals: List of PerformanceControlSignal objects to use for
        conditioning, or None if not conditioning on anything.
    optional_conditioning: If True, conditioning can be disabled by setting a
        flag as part of the conditioning input.
  """

  def __init__(self, details, encoder_decoder, hparams, num_velocity_bins=0,
               control_signals=None, optional_conditioning=False,
               note_performance=False):
    if control_signals is not None:
      control_encoder = note_seq.MultipleEventSequenceEncoder(
          [control.encoder for control in control_signals])
      if optional_conditioning:
        control_encoder = note_seq.OptionalEventSequenceEncoder(control_encoder)
      encoder_decoder = note_seq.ConditionalEventSequenceEncoderDecoder(
          control_encoder, encoder_decoder)

    super(PerformanceRnnConfig, self).__init__(
        details, encoder_decoder, hparams)
    self.num_velocity_bins = num_velocity_bins
    self.control_signals = control_signals
    self.optional_conditioning = optional_conditioning
    self.note_performance = note_performance


default_configs = {
    'performance':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='performance', description='Performance RNN'),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding()),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001)),
    'performance_with_dynamics':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='performance_with_dynamics',
                description='Performance RNN with dynamics'),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32),
    'performance_with_dynamics_compact':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='performance_with_dynamics',
                description='Performance RNN with dynamics (compact input)'),
            note_seq.OneHotIndexEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32),
    'performance_with_dynamics_and_modulo_encoding':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='performance_with_dynamics_and_modulo_encoding',
                description='Performance RNN with dynamics and modulo encoding'
            ),
            note_seq.ModuloPerformanceEventSequenceEncoderDecoder(
                num_velocity_bins=32),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32),
    'performance_with_dynamics_and_note_encoding':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='performance_with_dynamics_and_note_encoding',
                description='Performance RNN with dynamics and note encoding'),
            note_seq.NotePerformanceEventSequenceEncoderDecoder(
                num_velocity_bins=32),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32,
            note_performance=True),
    'density_conditioned_performance_with_dynamics':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='density_conditioned_performance_with_dynamics',
                description=
                'Note-density-conditioned Performance RNN + dynamics'
            ),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32,
            control_signals=[
                note_seq.NoteDensityPerformanceControlSignal(
                    window_size_seconds=3.0,
                    density_bin_ranges=[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0])
            ]),
    'pitch_conditioned_performance_with_dynamics':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='pitch_conditioned_performance_with_dynamics',
                description='Pitch-histogram-conditioned Performance RNN'),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32,
            control_signals=[
                note_seq.PitchHistogramPerformanceControlSignal(
                    window_size_seconds=5.0)
            ]),
    'multiconditioned_performance_with_dynamics':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='multiconditioned_performance_with_dynamics',
                description='Density- and pitch-conditioned Performance RNN'),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32,
            control_signals=[
                note_seq.NoteDensityPerformanceControlSignal(
                    window_size_seconds=3.0,
                    density_bin_ranges=[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0]),
                note_seq.PitchHistogramPerformanceControlSignal(
                    window_size_seconds=5.0)
            ]),
    'optional_multiconditioned_performance_with_dynamics':
        PerformanceRnnConfig(
            generator_pb2.GeneratorDetails(
                id='optional_multiconditioned_performance_with_dynamics',
                description='Optionally multiconditioned Performance RNN'),
            note_seq.OneHotEventSequenceEncoderDecoder(
                note_seq.PerformanceOneHotEncoding(num_velocity_bins=32)),
            contrib_training.HParams(
                batch_size=64,
                rnn_layer_sizes=[512, 512, 512],
                dropout_keep_prob=1.0,
                clip_norm=3,
                learning_rate=0.001),
            num_velocity_bins=32,
            control_signals=[
                note_seq.NoteDensityPerformanceControlSignal(
                    window_size_seconds=3.0,
                    density_bin_ranges=[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0]),
                note_seq.PitchHistogramPerformanceControlSignal(
                    window_size_seconds=5.0)
            ],
            optional_conditioning=True)
}
