product-vamp/generator/comping.py

"""Source-only accompaniment: joint layered rhythms and connected grips."""
from __future__ import annotations
import copy
import dataclasses
import functools
import json
import fractions
import typing

import generator.accompaniment
import generator.comping_rhythm
import generator.chord_families
import generator.acoustic_voicing
import generator.arrange.form
import generator.fretboard
import generator.guitar
import generator.harmony
import generator.model
import generator.parts
import generator.score
import generator.vocabulary
import generator.solo.policy

def semantic_key(chord):
  return json.dumps({k: chord[k] for k in ('root', 'kind', 'degrees', 'bass', 'inversion') if k in chord}, sort_keys=True)

def requirements(chord, part):
  result = generator.fretboard.chord_requirements(chord, respect_bass=part != 'electric')
  if part == 'electric' or result.basses:
    return result
  degrees = generator.harmony.degrees(chord)
  plain = degrees in ({1: 0, 3: 0, 5: 0}, {1: 0, 3: -1, 5: 0})
  root = int(generator.model.NATURAL[chord['root']['letter']] + generator.model.alteration(chord['root']['alter']))
  basses = tuple(sorted({(root + generator.harmony.degree_semitones(d, degrees[d])) % 12 for d in ((1, 3, 5) if plain else (1, 5)) if d in degrees}))
  if not basses:
    raise ValueError('acoustic accompaniment requires a written bass, root or fifth')
  return dataclasses.replace(result, basses=basses)

@functools.lru_cache(maxsize=2048)
def voicings(key, part):
  chord = json.loads(key)
  required = requirements(chord, part)
  if part == 'electric':
    options = generator.guitar.chord_positions(required, prune=False, profile='electric')
  elif part == 'acoustic':
    degrees = generator.harmony.degrees(chord)
    plain = degrees in ({1: 0, 3: 0, 5: 0}, {1: 0, 3: -1, 5: 0})
    options = tuple(v for v in generator.acoustic_voicing.candidates(required, plain)
      if v.bass.fret <= generator.fretboard.MAX_SINGLE_FRET
      and generator.fretboard.has_fret_anchor(v[1:])
      and (not plain or set(required.allowed) <= {p.midi % 12 for p in v}))
  else:
    raise ValueError('Unknown accompaniment instrument: ' + part)
  options = tuple(v for v in options if generator.fretboard.playable_attack((p, fractions.Fraction(1)) for p in v))
  if not options:
    layout = 'bass on 5/6 plus upper strings 2/3/4' if part == 'acoustic' else 'melody chord families on strings 1–4'
    raise ValueError(f'{generator.harmony.symbol(chord)}: no legal {part} voicing ({layout}); required degrees/bass, no open strings, fret anchor 1–12 and at most {generator.fretboard.MAX_FRET_SPAN} consecutive frets remain binding')
  return tuple(sorted(options))

def movement(part, previous, current):
  if not previous or not current:
    return 0
  if part == 'acoustic':
    return generator.acoustic_voicing.movement(previous, current)
  return generator.guitar.movement(previous, current)

class VoicingCost(typing.NamedTuple):
  extended_grips: int = 0
  family: int = 0
  effort: fractions.Fraction = fractions.Fraction()
  missing_tones: int = 0
  register: fractions.Fraction = fractions.Fraction()
  def plus(self, other):
    return VoicingCost(*(a + b for a, b in zip(self, other)))
  def moved(self, distance):
    return self._replace(effort=self.effort + distance)

def acoustic_preferences(grip, center):
  """Soft register/hand-position costs, in the same units as finger travel."""
  policy = generator.vocabulary.role('acoustic_accompaniment').get('voicing', {})
  register_weight = fractions.Fraction(str(policy.get('register_weight', 1)))
  position_weight = fractions.Fraction(str(policy.get('high_position_weight', 1)))
  threshold = fractions.Fraction(str(policy.get('comfortable_max_fret', 9)))
  if register_weight < 0 or position_weight < 0 or threshold < 1:
    raise ValueError('acoustic voicing requires nonnegative weights and a positive fret threshold')
  register = abs(fractions.Fraction(sum(p.midi for p in grip), len(grip)) - center)
  position = max(fractions.Fraction(), fractions.Fraction(sum(p.fret for p in grip), len(grip)) - threshold)
  return register_weight * register, position_weight * position

def local_cost(part, key, grip, center):
  required = requirements(json.loads(key), part)
  register = abs(fractions.Fraction(sum(p.midi for p in grip), len(grip)) - center)
  if part == 'electric':
    value = generator.guitar.voicing_cost(grip, 'electric', required.allowed)
    return VoicingCost(value.extended_grips, 3 * value.string_shortfall + value.family_penalty, value.effort, value.missing_tones, register)
  # Familiar degree orders on either bass string, then balanced whole-form
  # movement, sounding register and hand position. Neither is an absolute ban.
  family = generator.chord_families.acoustic(json.loads(key), grip)
  effort = sum(acoustic_preferences(grip, center))
  return VoicingCost(generator.fretboard.span_penalty(grip), int(not family['preferred']), effort, len(set(required.allowed) - {p.midi % 12 for p in grip}), register)

def select_grips(layers, costs, part, seed, *, left=(), right=(), close=False):
  """Exact small-state dynamic programming; never prune by hand center.

  Electric effort includes its melody grip cost and hand-center travel.
  Acoustic effort balances register/position with travel of all four voices,
  including a bass-string change. Compact span and family remain higher tiers.
  Fixing the first grip permits an exact repeat-join cost without a beam.
  """
  if not layers:
    return (), VoicingCost()
  # Every accompaniment transition is feasible (travel is a cost). A wider
  # option cannot beat a primary grip at the same span, regardless of travel.
  # This exact dominance reduction avoids growing the cubic repeat search.
  retained = []
  for i, layer in enumerate(layers):
    tier = min(costs[i][v].extended_grips for v in layer)
    retained.append(tuple(v for v in layer if costs[i][v].extended_grips == tier))
  layers = tuple(retained)
  ranks = [{v: generator.solo.policy.stable_rank(seed, (i, v)) for v in layer} for i, layer in enumerate(layers)]
  winners = []
  for first in layers[0] if close else [None]:
    states = {left: (VoicingCost(), (), ())}
    for i, layer in enumerate(layers):
      following = {}
      for grip in [first] if i == 0 and first is not None else layer:
        local = costs[i][grip]
        choices = []
        for previous, (cost, ties, path) in states.items():
          total = cost.plus(local).moved(movement(part, previous, grip))
          choices.append((total, (*ties, ranks[i][grip]), (*path, grip)))
        following[grip] = min(choices)
      states = following
    for grip, (cost, ties, path) in states.items():
      travel = movement(part, grip, first if close else right)
      winners.append((cost.moved(travel), ties, path))
  cost, _, path = min(winners)
  return path, cost

@dataclasses.dataclass(frozen=True)
class HarmonyChoice:
  index: int
  start: fractions.Fraction
  end: fractions.Fraction
  chord: dict
  events: tuple
  figure: dict
  options: tuple
  @property
  def key(self):
    return (self.index, self.start, self.end, semantic_key(self.chord), self.events)

def split_measures(flat, inputs, part):
  """Project a selected sounding figure into notation; barlines create ties only."""
  output, fragments, original = [], {}, {}
  cursor = fractions.Fraction()
  for m, prior, index in inputs:
    length = generator.model.fraction(m['duration'])
    result = generator.parts.build_measure(m, f'comp_{part}_{index}', 'rest', part, prior)
    voices = []
    for voice in flat['voices']:
      notes = []
      for n in voice['notes']:
        start = generator.model.fraction(n['onset'])
        end = start + generator.model.fraction(n['duration'])
        a, b = max(cursor, start), min(cursor + length, end)
        if a >= b:
          continue
        for duration, written in generator.accompaniment.written_durations(b - a):
          pitches = [{k: v for k, v in p.items() if k not in ('tie_from', 'tie_to')} for p in n['pitches']]
          event = generator.model.make_event(f'comp_{part}_{index}_{voice["id"]}_{len(notes)}', a - cursor, duration, 'note' if pitches else 'rest', pitches, written)
          event['voice'] = voice['id']
          for old, new in zip(n['pitches'], event['pitches']):
            fragments.setdefault(old['id'], []).append(new)
            original[old['id']] = old
          notes.append(event)
          a += duration
      voices.append({'id': voice['id'], 'notes': notes})
    result['voices'] = voices
    for chord in result['harmonies']:
      chord.pop('chord_optional', None)
    output.append(result)
    cursor += length
  for pieces in fragments.values():
    for a, b in zip(pieces, pieces[1:]):
      a['tie_to'], b['tie_from'] = b['id'], a['id']
  for identifier, old in original.items():
    if old.get('tie_to') in fragments:
      a, b = fragments[identifier][-1], fragments[old['tie_to']][0]
      a['tie_to'], b['tie_from'] = b['id'], a['id']
  return tuple(output)

class Canonical:
  def __init__(self, document, part, seed):
    self.part, self.seed = part, seed
    self.policy = generator.vocabulary.role(part + '_accompaniment')
    self.cache, self._reports, self._rhythm_reports = {}, {}, {}
    self.measures = generator.score.source_staff(document)['measures']
    self.melody = generator.comping_rhythm.melody_contexts(document)
    self.first = generator.score.body_start_measure_index(self.measures)
    sections = generator.arrange.form.full_arrangement_sections(self.measures, part)
    self.closed = any(a.role == b.role == 'accompaniment' and a.source_measure_index + a.measure_count == len(self.measures)
      and b.source_measure_index == self.first for a, b in zip(sections, sections[1:]))
    carried = generator.accompaniment.carried_chords(self.measures)
    self.groups = []
    index = self.first
    while index < len(self.measures):
      size = 1
      if index + 1 < len(self.measures):
        lengths = [generator.model.fraction(m['duration']) for m in self.measures[index:index + 2]]
        if any(tuple(map(fractions.Fraction, p['lengths'])) == tuple(lengths) for p in generator.vocabulary.accompaniment(part)['patterns']):
          size = 2
      self.groups.append(tuple((self.measures[i], carried[i], i) for i in range(index, index + size)))
      index += size
    layers = [self.candidates(group) for group in self.groups]
    selected = generator.comping_rhythm.select(layers, seed)
    self.block_rhythms = {group[0][2]: rhythm for group, rhythm in zip(self.groups, selected)}
    prepared = [self.prepare(group, rhythm) for group, rhythm in zip(self.groups, selected)]
    active = tuple(span for spans in prepared for span in spans if span.events)
    path, self.path_cost = self.select(active, close=self.closed)
    self.base = dict(zip((s.key for s in active), path))
    self.order = tuple((s.index, s.start, v) for s, v in zip(active, path))
    self.rhythms = {}
    previous = path[-1] if self.closed and path else ()
    for group, rhythm, spans in zip(self.groups, selected, prepared):
      grips = tuple(self.base[s.key] for s in spans if s.events)
      self.cache[self.key(group)] = self.render(group, rhythm, spans, grips, previous, False)
      cursor = fractions.Fraction()
      for m, _, i in group:
        length = generator.model.fraction(m['duration'])
        self.rhythms[i] = dataclasses.replace(rhythm, events=tuple(dataclasses.replace(e, start=max(e.start, cursor) - cursor, end=min(e.end, cursor + length) - cursor, harmonic_time=e.harmonic_time - cursor if e.harmonic_time is not None else None) for e in rhythm.events if e.start < cursor + length and e.end > cursor))
        cursor += length
      if grips:
        previous = grips[-1]
  @property
  def report(self):
    return [row for key in sorted(self._reports, key=repr) for row in self._reports[key]]
  @property
  def rhythm_report(self):
    return [row for key in sorted(self._rhythm_reports, key=repr) for row in self._rhythm_reports[key]]
  @staticmethod
  def key(inputs, limit=None):
    return (tuple((i, str(generator.model.fraction(m['duration'])), tuple((str(a), str(b), semantic_key(c) if c else '') for a, b, c in generator.comping_rhythm.spans(m, p))) for m, p, i in inputs), str(limit))
  def candidates(self, inputs, limit=None):
    flat = generator.comping_rhythm.flatten(inputs)
    melody, cursor, lengths = [], fractions.Fraction(), []
    for m, _, i in inputs:
      melody.extend((a + cursor, b + cursor) for a, b in self.melody[i])
      lengths.append(generator.model.fraction(m['duration']))
      cursor += lengths[-1]
    return generator.comping_rhythm.candidates(flat, None, self.part, tuple(melody), self.seed + '/' + str(inputs[0][2]), lengths=tuple(lengths), limit=limit)
  def prepare(self, inputs, rhythm):
    flat = generator.comping_rhythm.flatten(inputs)
    spans = []
    offsets, cursor = [], fractions.Fraction()
    for m, _, i in inputs:
      offsets.append((cursor, i))
      cursor += generator.model.fraction(m['duration'])
    for start, end, chord in generator.comping_rhythm.spans(flat, None):
      if not chord:
        continue
      offset, index = next((a, i) for a, i in reversed(offsets) if a <= start)
      events = tuple(e for e in rhythm.events if start <= (e.harmonic_time if e.harmonic_time is not None else e.start) < end)
      try:
        options = voicings(semantic_key(chord), self.part)
      except ValueError as error:
        raise ValueError(f'source measure {inputs[index - inputs[0][2]][0]["number"]}, beat {start - offset + 1}: {error}') from error
      spans.append(HarmonyChoice(index, start, end, chord, events, dict(rhythm.families), options))
    return tuple(spans)
  def select(self, spans, *, left=(), right=(), close=False):
    layers = tuple(s.options for s in spans)
    costs = tuple({v: local_cost(self.part, semantic_key(s.chord), v, self.policy.get('center', 55)) for v in layer} for s, layer in zip(spans, layers))
    return select_grips(layers, costs, self.part, self.seed, left=left, right=right, close=close)
  def __call__(self, source, prior, index, *, sources=None, bounds=None, stop=None):
    group = next((g for g in self.groups if any(i == index for _, _, i in g)), ((source, prior, index),))
    sources = sources or self.measures
    carried = generator.accompaniment.carried_chords(sources)
    low, high = bounds or (self.first, len(sources))
    group = tuple((source if i == index else sources[i], prior if i == index else carried[i], i) for _, _, i in group if low <= i < high)
    limit = None
    if stop and any(i == stop[0] for _, _, i in group):
      limit = sum((generator.model.fraction(m['duration']) for m, _, i in group if i < stop[0]), fractions.Fraction()) + stop[1]
    key = self.key(group, limit)
    if key not in self.cache:
      choices = self.candidates(group, limit)
      start, end = group[0][2], group[-1][2]
      # Preserve the canonical whole figure where still legal under changed
      # effective harmony. Recompute only genuinely different source context.
      original = self.block_rhythms.get(start)
      matches = [r for r in choices if original and r.pattern_id == original.pattern_id]
      rhythm = matches[0] if matches else generator.comping_rhythm.select((choices,), self.seed, self.block_rhythms.get(start - 2), self.block_rhythms.get(end + 1))[0]
      spans = self.prepare(group, rhythm)
      before = [v for i, _, v in self.order if i < start]
      after = [v for i, _, v in self.order if i > end]
      left = before[-1] if before else self.order[-1][2] if self.closed and self.order else ()
      right = after[0] if after else self.order[0][2] if self.closed and self.order else ()
      selected, _ = self.select(tuple(s for s in spans if s.events), left=left, right=right)
      self.cache[key] = self.render(group, rhythm, spans, selected, left, True, limit)
    return copy.deepcopy(self.cache[key][index - group[0][2]])
  def render(self, inputs, rhythm, spans, selected, previous, variant, limit=None):
    flat = generator.comping_rhythm.flatten(inputs)
    length = generator.model.fraction(flat['duration'])
    index = inputs[0][2]
    attacks, rows = [], []
    choices = iter(selected)
    for span in spans:
      grip = next(choices) if span.events else ()
      bass = grip[0] if grip and self.part == 'acoustic' else None
      degrees = generator.harmony.degrees(span.chord)
      root = int(generator.model.NATURAL[span.chord['root']['letter']] + generator.model.alteration(span.chord['root']['alter']))
      rows.append({'source_index': span.index, 'onset': str(span.start), 'chord': generator.harmony.symbol(span.chord),
        'figure': span.figure, 'pattern_id': rhythm.pattern_id, 'tones': len(grip), 'bass': bool(bass), 'variant': variant,
        'family': (generator.chord_families.acoustic(span.chord, grip) if self.part == 'acoustic' else generator.chord_families.electric(grip)) if grip else None,
        'family_coverage': generator.chord_families.coverage(span.chord, span.options, self.part),
        'fret_span': generator.fretboard.fret_span(grip),
        'span_group': ('fallback' if generator.fretboard.span_penalty(grip) else 'primary') if grip else None,
        'span_candidates': generator.chord_families.span_coverage(span.options),
        'positions': [{'string': p.string_number, 'fret': p.fret, 'sounding_midi': p.midi} for p in grip],
        'bass_degree': next((d for d, a in degrees.items() if bass and (root + generator.harmony.degree_semitones(d, a)) % 12 == bass.midi % 12), None),
        'bass_string': bass.string_number if bass else None,
        'movement': movement(self.part, previous, grip), 'voicing_candidates': len(span.options)})
      if grip and self.part == 'acoustic':
        register, position = acoustic_preferences(grip, self.policy.get('center', 55))
        travel = movement(self.part, previous, grip)
        rows[-1]['voicing_cost'] = {'movement': travel, 'register': str(register), 'position': str(position), 'combined': str(travel + register + position)}
      for event in span.events:
        lanes = ['bass', 'upper'] if event.lane == 'combined' else [event.lane]
        for lane in lanes:
          positions = (grip[0],) if lane == 'bass' else grip[1:] if lane == 'upper' else grip
          attacks.append((lane, event.start, event.end, positions, span.chord))
      if grip:
        previous = grip
    key = self.key(inputs, limit)
    self._reports[key] = rows
    voices = []
    for lane in ['upper', 'bass'] if self.part == 'acoustic' else ['electric']:
      notes, cursor = [], fractions.Fraction()
      def emit(a, b, positions=(), chord=None):
        previous = None
        for duration, written in generator.accompaniment.written_durations(b - a):
          pitches = generator.fretboard.chord_pitches(chord, positions) if positions else []
          n = generator.model.make_event(f'comp_{self.part}_{index}_{lane}_{len(notes)}', a, duration, 'note' if pitches else 'rest', pitches, written)
          n['voice'] = lane
          if previous and pitches:
            for old, new in zip(previous['pitches'], n['pitches']):
              old['tie_to'], new['tie_from'] = new['id'], old['id']
          notes.append(n)
          previous = n
          a += duration
      for _, start, end, positions, chord in sorted((e for e in attacks if e[0] == lane), key=lambda e: e[1]):
        if cursor < start:
          emit(cursor, start)
        emit(start, end, positions, chord)
        cursor = end
      if cursor < length:
        emit(cursor, length)
      voices.append({'id': lane, 'notes': notes})
    result = split_measures({'voices': voices}, inputs, self.part)
    reports, cursor = [], fractions.Fraction()
    for m, _, i in inputs:
      duration = generator.model.fraction(m['duration'])
      reports.append({'source_index': i, 'pattern_id': rhythm.pattern_id, 'pattern_start': index,
        'families': dict(rhythm.families), 'contributions': dict(rhythm.contributions), 'selection_cost': rhythm.cost,
        'exceptions': rhythm.exceptions, 'source_melody': [(str(a), str(b)) for a, b in self.melody[i]],
        'events': [{'layer': e.lane, 'onset': str(max(e.start, cursor) - cursor), 'duration': str(min(e.end, cursor + duration) - max(e.start, cursor)),
          'continuation': e.start < cursor, 'harmonic_target': str(e.harmonic_time - cursor) if e.harmonic_time is not None else None} for e in rhythm.events if e.start < cursor + duration and e.end > cursor]})
      cursor += duration
    self._rhythm_reports[key] = reports
    return result