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

"""Event sequence RNN model."""

import collections
import copy
import functools

from magenta.common import beam_search
from magenta.common import state_util
from magenta.models.shared import events_rnn_graph
import magenta.music as mm
import numpy as np
from six.moves import range  # pylint: disable=redefined-builtin
import tensorflow as tf

# Model state when generating event sequences, consisting of the next inputs to
# feed the model, the current RNN state, the current control sequence (if
# applicable), and state for the current control sequence (if applicable).
ModelState = collections.namedtuple(
    'ModelState', ['inputs', 'rnn_state', 'control_events', 'control_state'])


class EventSequenceRnnModelError(Exception):
  pass


def _extend_control_events_default(control_events, events, state):
  """Default function for extending control event sequence.

  This function extends a control event sequence by duplicating the final event
  in the sequence.  The control event sequence will be extended to have length
  one longer than the generated event sequence.

  Args:
    control_events: The control event sequence to extend.
    events: The list of generated events.
    state: State maintained while generating, unused.

  Returns:
    The resulting state after extending the control sequence (in this case the
    state will be returned unmodified).
  """
  while len(control_events) <= len(events):
    control_events.append(control_events[-1])
  return state


class EventSequenceRnnModel(mm.BaseModel):
  """Class for RNN event sequence generation models.

  Currently this class only supports generation, of both event sequences and
  note sequences (via event sequences). Support for model training will be added
  at a later time.
  """

  def __init__(self, config):
    """Initialize the EventSequenceRnnModel.

    Args:
      config: An EventSequenceRnnConfig containing the encoder/decoder and
        HParams to use.
    """
    super(EventSequenceRnnModel, self).__init__()
    self._config = config

  def _build_graph_for_generation(self):
    events_rnn_graph.get_build_graph_fn('generate', self._config)()

  def _batch_size(self):
    """Extracts the batch size from the graph."""
    return self._session.graph.get_collection('inputs')[0].shape[0].value

  def _generate_step_for_batch(self, event_sequences, inputs, initial_state,
                               temperature):
    """Extends a batch of event sequences by a single step each.

    This method modifies the event sequences in place.

    Args:
      event_sequences: A list of event sequences, each of which is a Python
          list-like object. The list of event sequences should have length equal
          to `self._batch_size()`. These are extended by this method.
      inputs: A Python list of model inputs, with length equal to
          `self._batch_size()`.
      initial_state: A numpy array containing the initial RNN state, where
          `initial_state.shape[0]` is equal to `self._batch_size()`.
      temperature: The softmax temperature.

    Returns:
      final_state: The final RNN state, a numpy array the same size as
          `initial_state`.
      loglik: The log-likelihood of the chosen softmax value for each event
          sequence, a 1-D numpy array of length
          `self._batch_size()`. If `inputs` is a full-length inputs batch, the
          log-likelihood of each entire sequence up to and including the
          generated step will be computed and returned.
    """
    assert len(event_sequences) == self._batch_size()

    graph_inputs = self._session.graph.get_collection('inputs')[0]
    graph_initial_state = self._session.graph.get_collection('initial_state')
    graph_final_state = self._session.graph.get_collection('final_state')
    graph_softmax = self._session.graph.get_collection('softmax')[0]
    graph_temperature = self._session.graph.get_collection('temperature')

    feed_dict = {graph_inputs: inputs,
                 tuple(graph_initial_state): initial_state}
    # For backwards compatibility, we only try to pass temperature if the
    # placeholder exists in the graph.
    if graph_temperature:
      feed_dict[graph_temperature[0]] = temperature
    final_state, softmax = self._session.run(
        [graph_final_state, graph_softmax], feed_dict)

    if isinstance(softmax, list):
      if softmax[0].shape[1] > 1:
        softmaxes = []
        for beam in range(softmax[0].shape[0]):
          beam_softmaxes = []
          for event in range(softmax[0].shape[1] - 1):
            beam_softmaxes.append(
                [softmax[s][beam, event] for s in range(len(softmax))])
          softmaxes.append(beam_softmaxes)
        loglik = self._config.encoder_decoder.evaluate_log_likelihood(
            event_sequences, softmaxes)
      else:
        loglik = np.zeros(len(event_sequences))
    else:
      if softmax.shape[1] > 1:
        # The inputs batch is longer than a single step, so we also want to
        # compute the log-likelihood of the event sequences up until the step
        # we're generating.
        loglik = self._config.encoder_decoder.evaluate_log_likelihood(
            event_sequences, softmax[:, :-1, :])
      else:
        loglik = np.zeros(len(event_sequences))