product-vamp/generator/arrange/generate.py

from __future__ import annotations

import copy
import fractions
import pathlib

import generator.accompaniment
import generator.arrange.accompaniment
import generator.acoustic
import generator.acoustic_voicing
import generator.arrange.form
import generator.arrange.pickup
import generator.electric
import generator.fretboard
import generator.guitar
import generator.melody
import generator.model
import generator.musicxml
import generator.parts
import generator.rhythm
import generator.scales.library
import generator.scales.mapping
import generator.scales.timeline
import generator.score
import generator.solo.arrangement
import generator.solo.sequences

FINAL_SECTION_INDEX = 6

def arrange_document(document: dict, part_type: str, *, rhythms: generator.rhythm.Library, scales: generator.scales.library.Library | None = None) -> dict:
  """Place the regular form; the song pipeline inserts solos after all heads exist."""
  document = generator.parts.prepare_document(document)
  plans = plan_arrangement(document, rhythms) if part_type != 'vocal' else {}
  return _arrange_document(document, part_type, plans, scales=scales)

def section_source(source: list[dict], section, next_section=None) -> list[dict]:
  result = copy.deepcopy(source)
  for measure in result:
    measure['harmonies'] = [h for h in measure['harmonies'] if section.section_index != FINAL_SECTION_INDEX or not h.get('chord_optional')]
    for chord in measure['harmonies']:
      chord.pop('chord_optional', None)
  if generator.arrange.pickup.enters(next_section, source):
    destination = result[section.source_measure_index + section.measure_count - 1]
    generator.arrange.pickup.harmonies(destination, result[0])
  return result

def plan_arrangement(document: dict, rhythms: generator.rhythm.Library) -> dict[int, generator.accompaniment.Plan]:
  source = generator.score.source_staff(document)['measures']
  body = generator.score.body_start_measure_index(source)
  previous = generator.accompaniment.carried_chords(source)
  canonical = generator.accompaniment.plan_section(source[body:], rhythms, initial=previous[body], electric=True)
  result = {}
  for section in generator.arrange.form.full_arrangement_sections(source, 'acoustic'):
    start = section.source_measure_index - body
    end = start + section.measure_count
    result[section.section_index] = generator.accompaniment.Plan(canonical.acoustic_bass[start:end],
      canonical.acoustic_upper[start:end], canonical.electric[start:end])
  return result

def _arrange_document(document: dict, part_type: str, plans: dict[int, generator.accompaniment.Plan], *, scales: generator.scales.library.Library | None = None, reports: list[str] | None = None, source_name: str = '', owned=None, accompaniment=None) -> dict:
  if part_type not in generator.parts.PART_TYPES:
    raise ValueError(f'unknown arrangement part: {part_type}')
  staff = generator.score.source_staff(document)
  source = staff['measures']
  sections = generator.arrange.form.full_arrangement_sections(source, part_type)
  output, tuplets, octave_shifts, voicing_origins = [], [], [], []
  solos = []
  absolute = fractions.Fraction()
  source_starts = []
  source_time = fractions.Fraction()
  for m in source:
    source_starts.append(source_time)
    source_time += generator.model.fraction(m['duration'])
  for index, section in enumerate(sections):
    role = section.role if owned is None or section.role in owned else 'rest'
    section_measures = section_source(source, section, sections[index + 1] if index + 1 < len(sections) else None)
    previous = generator.accompaniment.carried_chords(section_measures)
    segment = []
    for si in range(section.source_measure_index, section.source_measure_index + section.measure_count):
      m = section_measures[si]
      rhythm = None
      if role == 'accompaniment' and accompaniment is None:
        plan = plans[section.section_index]
        rhythm = (plan.acoustic_bass[len(segment)], plan.acoustic_upper[len(segment)]) if part_type == 'acoustic' else plan.electric[len(segment)]
      if role == 'accompaniment' and accompaniment is not None:
        result = accompaniment(m, previous[si], si)
        result['id'] = f'measure_{len(output) + len(segment) + 1}'
        result['section_label'] = section.label if not segment else ''
      else:
        result = generator.parts.build_measure(m, f'measure_{len(output) + len(segment) + 1}', role, part_type, previous[si], section.label if not segment else '', rhythm)
      segment.append(result)
      if part_type == 'electric' and section.role in ('melody', 'accompaniment'):
        voicing_origins.append(generator.electric.MeasureOrigin(source_starts[si], section.role, result))
    segment_length = sum((generator.model.fraction(m['duration']) for m in segment), fractions.Fraction())
    if section.role == 'slash':
      solos.append(generator.scales.timeline.Solo(absolute, absolute + segment_length, section.label))
    if role in ('vocal', 'melody'):
      start = source_starts[section.source_measure_index]
      for shift in staff.get('octave_shifts', []):
        low, high = max(start, generator.model.fraction(shift['start'])), min(start + segment_length, generator.model.fraction(shift['end']))
        if low < high:
          octave_shifts.append({**shift, 'id': f'section_{section.section_index}_{shift["id"]}', 'start': generator.model.rational(absolute + low - start), 'end': generator.model.rational(absolute + high - start)})
    pickup = None
    if role in ('vocal', 'melody') and generator.arrange.pickup.enters(section, source) and output:
      pickup = generator.parts.build_measure(source[0], f'pickup_{section.section_index}', section.role, part_type, None)
    tuplets.extend(generator.score.reidentify(([pickup] if pickup else []) + segment, staff.get('tuplets', [])))
    if pickup:
      destination = output[-1]
      pad = generator.arrange.pickup.offset(destination, pickup)
      onset = generator.arrange.pickup.splice(destination, pickup)
      if part_type == 'electric':
        voicing_origins = [generator.electric.MeasureOrigin(o.onset, o.role, o.measure, o.start, onset if o.measure is destination else o.end) for o in voicing_origins]
        voicing_origins.append(generator.electric.MeasureOrigin(-pad, 'melody', destination, onset))
      for shift in staff.get('octave_shifts', []):
        low, high = max(onset - pad, generator.model.fraction(shift['start'])), min(generator.model.fraction(pickup['duration']), generator.model.fraction(shift['end']))
        if low < high:
          offset = absolute - generator.model.fraction(pickup['duration'])
          octave_shifts.append({**shift, 'id': f'pickup_{section.section_index}_{shift["id"]}', 'start': generator.model.rational(offset + low), 'end': generator.model.rational(offset + high)})
    output.extend(segment)
    absolute += segment_length
  tuplets = generator.arrange.pickup.retain_links(output, tuplets)
  result = generator.score.make_arranged_document(document, staff, staff_id='staff_arranged_' + part_type, staff_name=part_type.capitalize(), role='vocal' if part_type == 'vocal' else 'guitar', arrangement_part_type=part_type, measures=output, tuplets=tuplets, octave_shifts=octave_shifts)
  result['metadata'].setdefault('tempo', 124)
  if part_type == 'acoustic' and accompaniment is None and (owned is None or 'accompaniment' in owned):
    generator.acoustic.prepare(result)
    generator.acoustic_voicing.place(result)
    generator.acoustic.combine(result)
  elif part_type == 'electric' and accompaniment is None and (owned is None or 'melody' in owned or 'accompaniment' in owned):
    generator.melody.harmonize(result)
    generator.guitar.place_guitar(result)
    if generator.accompaniment.connect_continuations(result):
      generator.guitar.place_guitar(result)
  # Source labels identify the head while arranging; published labels identify
  # the complete performance, including its lead-in and repeated choruses.
  for number, measure in enumerate(output, 1):
    measure['number'] = str(number)
  if part_type != 'vocal' and owned is None:
    generator.arrange.accompaniment.install(document, result, plans)
  generator.fretboard.validate_document(result)
  if scales and solos:
    report = generator.scales.mapping.annotate(result, solos, scales, source_name or document['metadata']['title'])
    if reports is not None:
      reports.append(report)
  return result

def generate_performance_scores(source_dir: pathlib.Path, output_dir: pathlib.Path, *, rhythms=None, scales=None, sequences=None, composition_jobs=2, cache_dir=None, progress=False) -> int:
  from generator.production import generate_tree
  return generate_tree(source_dir, output_dir, jobs=composition_jobs, cache_dir=cache_dir, progress=progress)