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

"""A module for implementing interaction between MIDI and SequenceGenerators."""

import abc
import threading
import time

import note_seq
from note_seq.protobuf import generator_pb2
from note_seq.protobuf import music_pb2
import tensorflow.compat.v1 as tf


def adjust_sequence_times(sequence, delta_time):
  """Adjusts note and total NoteSequence times by `delta_time`."""
  retimed_sequence = music_pb2.NoteSequence()
  retimed_sequence.CopyFrom(sequence)

  for note in retimed_sequence.notes:
    note.start_time += delta_time
    note.end_time += delta_time
  retimed_sequence.total_time += delta_time
  return retimed_sequence


class MidiInteraction(threading.Thread):
  """Base class for handling interaction between MIDI and SequenceGenerator.

  Child classes will provided the "main loop" of an interactive session between
  a MidiHub used for MIDI I/O and sequences generated by a SequenceGenerator in
  their `run` methods.

  Should be started by calling `start` to launch in a separate thread.

  Args:
    midi_hub: The MidiHub to use for MIDI I/O.
    sequence_generators: A collection of SequenceGenerator objects.
    qpm: The quarters per minute to use for this interaction. May be overriden
       by control changes sent to `tempo_control_number`.
    generator_select_control_number: An optional MIDI control number whose
       value to use for selection a sequence generator from the collection.
       Must be provided if `sequence_generators` contains multiple
       SequenceGenerators.
    tempo_control_number: An optional MIDI control number whose value to use to
       determine the qpm for this interaction. On receipt of a control change,
       the qpm will be set to 60 more than the control change value.
    temperature_control_number: The optional control change number to use for
        controlling generation softmax temperature.

  Raises:
    ValueError: If `generator_select_control_number` is None and
        `sequence_generators` contains multiple SequenceGenerators.
  """
  _metaclass__ = abc.ABCMeta

  # Base QPM when set by a tempo control change.
  _BASE_QPM = 60

  def __init__(self,
               midi_hub,
               sequence_generators,
               qpm,
               generator_select_control_number=None,
               tempo_control_number=None,
               temperature_control_number=None):
    if generator_select_control_number is None and len(sequence_generators) > 1:
      raise ValueError(
          '`generator_select_control_number` cannot be None if there are '
          'multiple SequenceGenerators.')
    self._midi_hub = midi_hub
    self._sequence_generators = sequence_generators
    self._default_qpm = qpm
    self._generator_select_control_number = generator_select_control_number
    self._tempo_control_number = tempo_control_number
    self._temperature_control_number = temperature_control_number

    # A signal to tell the main loop when to stop.
    self._stop_signal = threading.Event()
    super(MidiInteraction, self).__init__()

  @property
  def _sequence_generator(self):
    """Returns the SequenceGenerator selected by the current control value."""
    if len(self._sequence_generators) == 1:
      return self._sequence_generators[0]
    val = self._midi_hub.control_value(self._generator_select_control_number)
    val = 0 if val is None else val
    return self._sequence_generators[val % len(self._sequence_generators)]

  @property
  def _qpm(self):
    """Returns the qpm based on the current tempo control value."""
    val = self._midi_hub.control_value(self._tempo_control_number)
    return self._default_qpm if val is None else val + self._BASE_QPM

  def _temperature(self, min_temp=0.1, max_temp=2.0, default=1.0):
    """Returns the temperature based on the current control value.

    Linearly interpolates between `min_temp` and `max_temp`.

    Args:
      min_temp: The minimum temperature, which will be returned when value is 0.
      max_temp: The maximum temperature, which will be returned when value is
          127.
      default: The temperature to return if control value is None.

    Returns:
      A float temperature value based on the 8-bit MIDI control value.
    """
    val = self._midi_hub.control_value(self._temperature_control_number)
    if val is None:
      return default
    return min_temp + (val / 127.) * (max_temp - min_temp)

  @abc.abstractmethod
  def run(self):
    """The main loop for the interaction.

    Must exit shortly after `self._stop_signal` is set.
    """
    pass

  def stop(self):
    """Stops the main loop, and blocks until the interaction is stopped."""
    self._stop_signal.set()
    self.join()


class CallAndResponseMidiInteraction(MidiInteraction):
  """Implementation of a MidiInteraction for interactive "call and response".

  Alternates between receiving input from the MidiHub ("call") and playing
  generated sequences ("response"). During the call stage, the input is captured
  and used to generate the response, which is then played back during the
  response stage.

  The call phrase is started when notes are received and ended by an external
  signal (`end_call_signal`) or after receiving no note events for a full tick.
  The response phrase is immediately generated and played. Its length is
  optionally determined by a control value set for
  `response_ticks_control_number` or by the length of the call.

  Args:
    midi_hub: The MidiHub to use for MIDI I/O.
    sequence_generators: A collection of SequenceGenerator objects.
    qpm: The quarters per minute to use for this interaction. May be overriden
       by control changes sent to `tempo_control_number`.
    generator_select_control_number: An optional MIDI control number whose
       value to use for selection a sequence generator from the collection.
       Must be provided if `sequence_generators` contains multiple
       SequenceGenerators.
    clock_signal: An optional midi_hub.MidiSignal to use as a clock. Each tick
        period should have the same duration. No other assumptions are made
        about the duration, but is typically equivalent to a bar length. Either
        this or `tick_duration` must be specified.be
    tick_duration: An optional float specifying the duration of a tick period in
        seconds. No assumptions are made about the duration, but is typically
        equivalent to a bar length. Either this or `clock_signal` must be
        specified.
    end_call_signal: The optional midi_hub.MidiSignal to use as a signal to stop
        the call phrase at the end of the current tick.
    panic_signal: The optional midi_hub.MidiSignal to use as a signal to end
        all open notes and clear the playback sequence.
    mutate_signal: The optional midi_hub.MidiSignal to use as a signal to
        generate a new response sequence using the current response as the
        input.
    allow_overlap: A boolean specifying whether to allow the call to overlap
        with the response.
    metronome_channel: The optional 0-based MIDI channel to output metronome on.
        Ignored if `clock_signal` is provided.
    min_listen_ticks_control_number: The optional control change number to use
        for controlling the minimum call phrase length in clock ticks.
    max_listen_ticks_control_number: The optional control change number to use
        for controlling the maximum call phrase length in clock ticks. Call
        phrases will automatically be ended and responses generated when this
        length is reached.
    response_ticks_control_number: The optional control change number to use for
        controlling the length of the response in clock ticks.
    tempo_control_number: An optional MIDI control number whose value to use to
       determine the qpm for this interaction. On receipt of a control change,
       the qpm will be set to 60 more than the control change value.
    temperature_control_number: The optional control change number to use for
        controlling generation softmax temperature.
    loop_control_number: The optional control change number to use for
        determining whether the response should be looped. Looping is enabled
        when the value is 127 and disabled otherwise.
    state_control_number: The optinal control change number to use for sending
        state update control changes. The values are 0 for `IDLE`, 1 for
        `LISTENING`, and 2 for `RESPONDING`.

    Raises:
      ValueError: If exactly one of `clock_signal` or `tick_duration` is not
         specified.
  """

  class State(object):
    """Class holding state value representations."""
    IDLE = 0
    LISTENING = 1
    RESPONDING = 2

    _STATE_NAMES = {
        IDLE: 'Idle', LISTENING: 'Listening', RESPONDING: 'Responding'}

    @classmethod
    def to_string(cls, state):
      return cls._STATE_NAMES[state]

  def __init__(self,
               midi_hub,
               sequence_generators,
               qpm,
               generator_select_control_number,
               clock_signal=None,
               tick_duration=None,
               end_call_signal=None,
               panic_signal=None,
               mutate_signal=None,
               allow_overlap=False,
               metronome_channel=None,
               min_listen_ticks_control_number=None,
               max_listen_ticks_control_number=None,
               response_ticks_control_number=None,
               tempo_control_number=None,
               temperature_control_number=None,
               loop_control_number=None,
               state_control_number=None):
    super(CallAndResponseMidiInteraction, self).__init__(
        midi_hub, sequence_generators, qpm, generator_select_control_number,
        tempo_control_number, temperature_control_number)
    if [clock_signal, tick_duration].count(None) != 1:
      raise ValueError(
          'Exactly one of `clock_signal` or `tick_duration` must be specified.')
    self._clock_signal = clock_signal
    self._tick_duration = tick_duration