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

"""Utility functions for working with chord progressions.

Use extract_chords_for_melodies to extract chord progressions from a
quantized NoteSequence object, aligned with already-extracted melodies.

Use ChordProgression.to_sequence to write a chord progression to a
NoteSequence proto, encoding the chords as text annotations.
"""

import abc
import bisect

from note_seq import chord_symbols_lib
from note_seq import constants
from note_seq import events_lib
from note_seq import sequences_lib
from note_seq.protobuf import music_pb2

STANDARD_PPQ = constants.STANDARD_PPQ
NOTES_PER_OCTAVE = constants.NOTES_PER_OCTAVE
NO_CHORD = constants.NO_CHORD

# Shortcut to CHORD_SYMBOL annotation type.
CHORD_SYMBOL = music_pb2.NoteSequence.TextAnnotation.CHORD_SYMBOL


class CoincidentChordsError(Exception):
  pass


class BadChordError(Exception):
  pass


class ChordProgression(events_lib.SimpleEventSequence):
  """Stores a quantized stream of chord events.

  ChordProgression is an intermediate representation that all chord or lead
  sheet models can use. Chords are represented here by a chord symbol string;
  model-specific code is responsible for converting this representation to
  SequenceExample protos for TensorFlow.

  ChordProgression implements an iterable object. Simply iterate to retrieve
  the chord events.

  ChordProgression events are chord symbol strings like "Cm7", with special
  event NO_CHORD to indicate no chordal harmony. When a chord lasts for longer
  than a single step, the chord symbol event is repeated multiple times. Note
  that this is different from Melody, where the special MELODY_NO_EVENT is used
  for subsequent steps of sustained notes; in the case of harmony, there's no
  distinction between a repeated chord and a sustained chord.

  Chords must be inserted in ascending order by start time.

  Attributes:
    start_step: The offset of the first step of the progression relative to the
        beginning of the source sequence.
    end_step: The offset to the beginning of the bar following the last step
       of the progression relative to the beginning of the source sequence.
    steps_per_quarter: Number of steps in in a quarter note.
    steps_per_bar: Number of steps in a bar (measure) of music.
  """

  def __init__(self, events=None, **kwargs):
    """Construct a ChordProgression."""
    if 'pad_event' in kwargs:
      del kwargs['pad_event']
    super(ChordProgression, self).__init__(pad_event=NO_CHORD,
                                           events=events, **kwargs)

  def _add_chord(self, figure, start_step, end_step):
    """Adds the given chord to the `events` list.

    `start_step` is set to the given chord. Everything after `start_step` in
    `events` is deleted before the chord is added. `events`'s length will be
     changed so that the last event has index `end_step` - 1.

    Args:
      figure: Chord symbol figure. A string like "Cm9" representing the chord.
      start_step: A non-negative integer step that the chord begins on.
      end_step: An integer step that the chord ends on. The chord is considered
          to end at the onset of the end step. `end_step` must be greater than
          `start_step`.

    Raises:
      BadChordError: If `start_step` does not precede `end_step`.
    """
    if start_step >= end_step:
      raise BadChordError(
          'Start step does not precede end step: start=%d, end=%d' %
          (start_step, end_step))

    self.set_length(end_step)

    for i in range(start_step, end_step):
      self._events[i] = figure

  def from_quantized_sequence(self, quantized_sequence, start_step, end_step):
    """Populate self with the chords from the given quantized NoteSequence.

    A chord progression is extracted from the given sequence starting at time
    step `start_step` and ending at time step `end_step`.

    The number of time steps per bar is computed from the time signature in
    `quantized_sequence`.

    Args:
      quantized_sequence: A quantized NoteSequence instance.
      start_step: Start populating chords at this time step.
      end_step: Stop populating chords at this time step.

    Raises:
      NonIntegerStepsPerBarError: If `quantized_sequence`'s bar length
          (derived from its time signature) is not an integer number of time
          steps.
      CoincidentChordsError: If any of the chords start on the same step.
    """
    sequences_lib.assert_is_relative_quantized_sequence(quantized_sequence)
    self._reset()

    steps_per_bar_float = sequences_lib.steps_per_bar_in_quantized_sequence(
        quantized_sequence)
    if steps_per_bar_float % 1 != 0:
      raise events_lib.NonIntegerStepsPerBarError(
          'There are %f timesteps per bar. Time signature: %d/%d' %
          (steps_per_bar_float, quantized_sequence.time_signature.numerator,
           quantized_sequence.time_signature.denominator))
    self._steps_per_bar = int(steps_per_bar_float)
    self._steps_per_quarter = (
        quantized_sequence.quantization_info.steps_per_quarter)

    # Sort track by chord times.
    chords = sorted([a for a in quantized_sequence.text_annotations
                     if a.annotation_type == CHORD_SYMBOL],
                    key=lambda chord: chord.quantized_step)

    prev_step = None
    prev_figure = NO_CHORD

    for chord in chords:
      if chord.quantized_step >= end_step:
        # No more chords within range.
        break

      elif chord.quantized_step < start_step:
        # Chord is before start of range.
        prev_step = chord.quantized_step
        prev_figure = chord.text
        continue

      if chord.quantized_step == prev_step:
        if chord.text == prev_figure:
          # Identical coincident chords, just skip.
          continue
        else:
          # Two different chords start at the same time step.
          self._reset()
          raise CoincidentChordsError(
              'chords %s and %s are coincident' % (prev_figure, chord.text))

      if chord.quantized_step > start_step:
        # Add the previous chord.
        if prev_step is None:
          start_index = 0
        else:
          start_index = max(prev_step, start_step) - start_step
        end_index = chord.quantized_step - start_step
        self._add_chord(prev_figure, start_index, end_index)

      prev_step = chord.quantized_step
      prev_figure = chord.text

    if prev_step is None or prev_step < end_step:
      # Add the last chord active before end_step.
      if prev_step is None:
        start_index = 0
      else:
        start_index = max(prev_step, start_step) - start_step
      end_index = end_step - start_step
      self._add_chord(prev_figure, start_index, end_index)

    self._start_step = start_step
    self._end_step = end_step

  def to_sequence(self,
                  sequence_start_time=0.0,
                  qpm=120.0):
    """Converts the ChordProgression to NoteSequence proto.

    This doesn't generate actual notes, but text annotations specifying the
    chord changes when they occur.

    Args:
      sequence_start_time: A time in seconds (float) that the first chord in
          the sequence will land on.
      qpm: Quarter notes per minute (float).

    Returns:
      A NoteSequence proto encoding the given chords as text annotations.
    """
    seconds_per_step = 60.0 / qpm / self.steps_per_quarter

    sequence = music_pb2.NoteSequence()
    sequence.tempos.add().qpm = qpm
    sequence.ticks_per_quarter = STANDARD_PPQ

    current_figure = NO_CHORD
    for step, figure in enumerate(self):
      if figure != current_figure:
        current_figure = figure
        chord = sequence.text_annotations.add()
        chord.time = step * seconds_per_step + sequence_start_time
        chord.text = figure
        chord.annotation_type = CHORD_SYMBOL

    return sequence

  def transpose(self, transpose_amount):
    """Transpose chords in this ChordProgression.

    Args:
      transpose_amount: The number of half steps to transpose this
          ChordProgression. Positive values transpose up. Negative values
          transpose down.

    Raises:
      ChordSymbolError: If a chord (other than "no chord") fails to be
          interpreted by the `chord_symbols_lib` module.
    """
    for i in range(len(self._events)):
      if self._events[i] != NO_CHORD:
        self._events[i] = chord_symbols_lib.transpose_chord_symbol(
            self._events[i], transpose_amount % NOTES_PER_OCTAVE)


def event_list_chords(quantized_sequence, event_lists):
  """Extract corresponding chords for multiple EventSequences.

  Args:
    quantized_sequence: The underlying quantized NoteSequence from which to
        extract the chords. It is assumed that the step numbering in this
        sequence matches the step numbering in each EventSequence in
        `event_lists`.
    event_lists: A list of EventSequence objects.

  Returns:
    A nested list of chord the same length as `event_lists`, where each list is
    the same length as the corresponding EventSequence (in events, not steps).
  """
  sequences_lib.assert_is_relative_quantized_sequence(quantized_sequence)

  chords = ChordProgression()
  if quantized_sequence.total_quantized_steps > 0:
    chords.from_quantized_sequence(
        quantized_sequence, 0, quantized_sequence.total_quantized_steps)

  pad_chord = chords[-1] if chords else NO_CHORD

  chord_lists = []
  for e in event_lists:
    chord_lists.append([chords[step] if step < len(chords) else pad_chord
                        for step in e.steps])

  return chord_lists


def event_list_keys(sequence, event_lists, steps_per_second):
  """Extract corresponding keys for multiple EventSequences.

  Args:
    sequence: The underlying NoteSequence from which to extract the keys.
    event_lists: A list of EventSequence objects.
    steps_per_second: The number of quantized steps per second in the event
        lists.

  Returns:
    A nested list of keys (integers 0-11 indicating the key signature) the same
    length as `event_lists`, where each list is the same length as the
    corresponding EventSequence (in events, not necessarily steps).
  """
  if not sequence.key_signatures:
    raise ValueError('Sequence has no key signatures.')

  key_changes = sorted((steps_per_second * ks.time, ks.key)
                       for ks in sequence.key_signatures)
  key_change_steps = [step for (step, _) in key_changes]

  def step_to_key(step):
    index = bisect.bisect(key_change_steps, step)
    if index == 0:
      # Step is before first key signature; just use that key.
      return key_changes[0][1]
    else:
      return key_changes[index - 1][1]

  key_lists = []
  for e in event_lists:
    key_lists.append([step_to_key(step) for step in e.steps])

  return key_lists


def add_chords_to_sequence(note_sequence, chords, chord_times):
  """Add chords to a NoteSequence (in place) at specified times.

  Args:
    note_sequence: The NoteSequence proto to which chords will be added (in
        place). Should not already have chords.
    chords: A Python list of chord figure strings to add to `note_sequence` as
        text annotations.
    chord_times: A Python list containing the time in seconds at which to add
        each chord. Should be the same length as `chords` and nondecreasing.

  Raises:
    ValueError: If `note_sequence` already has chords, or if `chord_times` is
        not sorted in ascending order.
  """
  if any(ta.annotation_type == CHORD_SYMBOL
         for ta in note_sequence.text_annotations):
    raise ValueError('NoteSequence already has chords.')
  if any(t1 > t2 for t1, t2 in zip(chord_times[:-1], chord_times[1:])):
    raise ValueError('Chord times not sorted in ascending order.')

  current_chord = None
  for chord, time in zip(chords, chord_times):
    if chord != current_chord:
      current_chord = chord
      ta = note_sequence.text_annotations.add()
      ta.annotation_type = CHORD_SYMBOL
      ta.time = time
      ta.text = chord


def add_keys_to_sequence(note_sequence, keys, key_times):
  """Add key signatures to a NoteSequence (in place) at specified times.

  Args:
    note_sequence: The NoteSequence proto to which key signatures will be added
        (in place). Should not already have key signatures.
    keys: A Python list of keys (integers 0-11) to add to `note_sequence` as
        KeySignature fields.
    key_times: A Python list containing the time in seconds at which to add
        each key signature. Should be the same length as `keys` and
        nondecreasing.

  Raises:
    ValueError: If `note_sequence` already has key signatures, or if `key_times`
        is not sorted in ascending order.
  """
  if note_sequence.key_signatures:
    raise ValueError('NoteSequence already has keys.')
  if any(t1 > t2 for t1, t2 in zip(key_times[:-1], key_times[1:])):
    raise ValueError('Key times not sorted in ascending order.')

  current_key = None
  for key, time in zip(keys, key_times):
    if key != current_key:
      current_key = key
      ks = note_sequence.key_signatures.add()
      ks.time = time
      ks.key = key


class ChordRenderer(object):
  """An abstract class for rendering NoteSequence chord symbols as notes."""
  __metaclass__ = abc.ABCMeta

  @abc.abstractmethod
  def render(self, sequence):
    """Renders the chord symbols of a NoteSequence.

    This function renders chord symbol annotations in a NoteSequence as actual
    notes. Notes are added to the NoteSequence object, and the chord symbols
    remain also.

    Args:
      sequence: The NoteSequence for which to render chord symbols.
    """
    pass


class BasicChordRenderer(ChordRenderer):
  """A chord renderer that holds each note for the duration of the chord."""

  def __init__(self,
               velocity=100,
               instrument=1,
               program=88,
               octave=4,
               bass_octave=3):
    """Initialize a BasicChordRenderer object.

    Args:
      velocity: The MIDI note velocity to use.
      instrument: The MIDI instrument to use.
      program: The MIDI program to use.
      octave: The octave in which to render chord notes. If the bass note is not
          otherwise part of the chord, it will not be rendered in this octave.
      bass_octave: The octave in which to render chord bass notes.
    """
    self._velocity = velocity
    self._instrument = instrument
    self._program = program
    self._octave = octave
    self._bass_octave = bass_octave

  def _render_notes(self, sequence, pitches, bass_pitch, start_time, end_time):
    """Renders notes."""
    all_pitches = []
    for pitch in pitches:
      all_pitches.append(12 * self._octave + pitch % 12)
    all_pitches.append(12 * self._bass_octave + bass_pitch % 12)

    for pitch in all_pitches:
      # Add a note.
      note = sequence.notes.add()
      note.start_time = start_time
      note.end_time = end_time
      note.pitch = pitch
      note.velocity = self._velocity
      note.instrument = self._instrument
      note.program = self._program

  def render(self, sequence):
    # Sort text annotations by time.
    annotations = sorted(sequence.text_annotations, key=lambda a: a.time)

    prev_time = 0.0
    prev_figure = NO_CHORD

    for annotation in annotations:
      if annotation.time >= sequence.total_time:
        break

      if annotation.annotation_type == CHORD_SYMBOL:
        if prev_figure != NO_CHORD:
          # Render the previous chord.
          pitches = chord_symbols_lib.chord_symbol_pitches(prev_figure)
          bass_pitch = chord_symbols_lib.chord_symbol_bass(prev_figure)
          self._render_notes(sequence=sequence,
                             pitches=pitches,
                             bass_pitch=bass_pitch,
                             start_time=prev_time,
                             end_time=annotation.time)

        prev_time = annotation.time
        prev_figure = annotation.text

    if (prev_time < sequence.total_time and
        prev_figure != NO_CHORD):
      # Render the last chord.
      pitches = chord_symbols_lib.chord_symbol_pitches(prev_figure)
      bass_pitch = chord_symbols_lib.chord_symbol_bass(prev_figure)
      self._render_notes(sequence=sequence,
                         pitches=pitches,
                         bass_pitch=bass_pitch,
                         start_time=prev_time,
                         end_time=sequence.total_time)