product-vamp/generator/engraving/notation.py

from __future__ import annotations

import bisect
import dataclasses
import fractions

import generator.engraving.fonts
import generator.engraving.types
import generator.model

ACCIDENTAL_GLYPHS = {fractions.Fraction(0): 'accidentalNatural', fractions.Fraction(1): 'accidentalSharp', fractions.Fraction(-1): 'accidentalFlat', fractions.Fraction(2): 'accidentalDoubleSharp', fractions.Fraction(-2): 'accidentalDoubleFlat', fractions.Fraction(3): 'accidentalTripleSharp', fractions.Fraction(-3): 'accidentalTripleFlat', fractions.Fraction(1, 2): 'accidentalQuarterToneSharpStein', fractions.Fraction(-1, 2): 'accidentalQuarterToneFlatStein', fractions.Fraction(3, 2): 'accidentalThreeQuarterTonesSharpStein', fractions.Fraction(-3, 2): 'accidentalThreeQuarterTonesFlatZimmermann'}
HEADS = {'x': 'X', 'diamond': 'Diamond', 'triangle': 'TriangleUp', 'inverted triangle': 'TriangleDown', 'circle-x': 'CircleX', 'cross': 'Plus', 'slash': 'Slash', 'normal': '', 'do': 'TriangleUp', 're': 'Moon', 'mi': 'Diamond', 'fa': 'TriangleRight', 'fa up': 'TriangleLeft', 'so': 'Round', 'la': 'Square', 'ti': 'TriangleRoundDown', 'rectangle': 'Square', 'square': 'Square', 'arrow down': 'LargeArrowDown', 'arrow up': 'LargeArrowUp'}

@dataclasses.dataclass(frozen=True)
class Event:
  measure: dict
  voice: str
  lane: int
  lanes: int
  start: fractions.Fraction
  end: fractions.Fraction
  value: dict

class ScoreIndex:
  def __init__(self, document: dict):
    self.staff = document['staves'][0]
    self.measures = []
    self.events = []
    self.events_by_onset = {}
    self.attributes_by_onset = {}
    self.harmonies_by_onset = {}
    self.by_id = {}
    self.tones = {}
    self.boundaries = {fractions.Fraction()}
    self.columns = {fractions.Fraction()}
    self._states = {}
    self._beam_groups = None
    time = fractions.Fraction()
    for measure in self.staff['measures']:
      end = time + generator.model.fraction(measure['duration'])
      self.measures.append((time, end, measure))
      self.boundaries.add(end)
      self.columns.update((time, end))
      for lane, voice in enumerate(measure['voices']):
        for note in voice['notes']:
          onset = time + generator.model.fraction(note['onset'])
          event = Event(measure, voice['id'], lane, len(measure['voices']), onset, onset + generator.model.fraction(note['duration']), note)
          self.events.append(event)
          self.by_id[note['id']] = event
          self.columns.add(onset)
          for tone in note['pitches']:
            self.tones[tone['id']] = (event, tone)
      for field, bucket in (('attributes', self.attributes_by_onset), ('harmonies', self.harmonies_by_onset)):
        for index, item in enumerate(measure.get(field, [])):
          onset = time + generator.model.fraction(item['onset'])
          self.columns.add(onset)
          bucket.setdefault(onset, []).append((measure, index, item))
      time = end
    self.end = time
    self.measure_starts = [item[0] for item in self.measures]
    self.measure_starts_by_id = {measure['id']: start for start, _, measure in self.measures}
    self.events.sort(key=lambda e: (e.start, not bool(e.value.get('grace')), e.lane))
    self.event_starts = []
    self.event_max_ends = []
    for event in self.events:
      self.events_by_onset.setdefault(event.start, []).append(event)
      self.event_starts.append(event.start)
      self.event_max_ends.append(max(event.end, self.event_max_ends[-1] if self.event_max_ends else event.end))
    self.onsets = list(self.events_by_onset)
    self.column_times = sorted(self.columns)
    self.boundary_times = sorted(self.boundaries)
    self.lyric_lanes = list(dict.fromkeys((e.voice, verse['number']) for e in self.events for verse in e.value.get('lyrics', [])))
    self.accidentals = self.accidental_state()
  def measure_at(self, time: fractions.Fraction, before: bool = False) -> tuple:
    position = bisect.bisect_left(self.measure_starts, time) if before else bisect.bisect_right(self.measure_starts, time)
    return self.measures[max(0, position - 1)]
  def events_at(self, time: fractions.Fraction, include_sustained: bool = False) -> list[Event]:
    attacks = self.events_by_onset.get(time, [])
    if not include_sustained:
      return attacks
    sustained = []
    position = bisect.bisect_left(self.event_starts, time) - 1
    # Prefix maxima let system continuations stop searching once every earlier note has ended.
    while position >= 0 and self.event_max_ends[position] > time:
      event = self.events[position]
      if event.end > time:
        sustained.append(event)
      position -= 1
    return list(reversed(sustained)) + attacks
  def state(self, time: fractions.Fraction, before: bool = False) -> dict:
    if (time, before) in self._states:
      return self._states[time, before]
    start, _, measure = self.measure_at(time, before)
    state = {key: measure[key] for key in ('key', 'meter', 'clef', 'transpose', 'staff_details')}
    for change in measure.get('attributes', []):
      onset = start + generator.model.fraction(change['onset'])
      if onset < time or (onset == time and not before):
        state.update({key: value for key, value in change.items() if key != 'onset'})
    self._states[time, before] = state
    return state
  def pitch_position(self, tone: dict, time: fractions.Fraction) -> float:
    state = self.state(time)
    clef = state['clef']
    pitch = tone.get('unpitched', tone)
    sign = clef['sign']
    reference = {'G': 4 * 7 + 4, 'F': 3 * 7 + 3, 'C': 4 * 7}.get(sign, 4 * 7 + 4)
    clef_line = clef.get('line', 2) if sign in ('G', 'F', 'C') else 2
    octave = 0 if 'unpitched' in tone else clef.get('octave_change', 0)
    shift = sum(s['octaves'] for s in self.staff.get('octave_shifts', []) if generator.model.fraction(s['start']) <= time < generator.model.fraction(s['end'])) if 'unpitched' not in tone else 0
    return pitch['octave'] * 7 + generator.model.LETTERS.index(pitch['letter']) + 7 * (shift - octave) - reference + 2 * (clef_line - 1)
  def accidental_state(self) -> dict[str, bool]:
    result = {}
    for start, end, measure in self.measures:
      memory = {}
      key = None
      onsets = self.onsets[bisect.bisect_left(self.onsets, start):bisect.bisect_left(self.onsets, end)]
      for onset in onsets:
        state = self.state(onset)
        if key != state['key']:
          key, memory = state['key'], {}
        simultaneous = {}
        for event in self.events_by_onset[onset]:
          for tone in event.value['pitches']:
            hidden_head = isinstance(tone.get('notehead'), dict) and tone['notehead'].get('name') == 'none'
            if event.value.get('hidden') or tone.get('print', {}).get('print-object') == 'no' or hidden_head:
              result[tone['id']] = False
              continue
            if 'unpitched' in tone:
              result[tone['id']] = False
              continue
            identity = (tone['letter'], tone['octave'])
            value = generator.model.alteration(tone['alter'])
            signature = fractions.Fraction() if key.get('custom') else generator.model.alteration(key['alterations'][generator.model.LETTERS.index(tone['letter'])])
            for step in key.get('steps', []):
              if step['letter'] == tone['letter'] and step.get('octave', tone['octave']) == tone['octave']:
                signature = generator.model.alteration(step['alter'])
            result[tone['id']] = bool(tone.get('accidental_display')) or (not tone.get('tie_from') and memory.get(identity, signature) != value)
            simultaneous.setdefault(identity, set()).add(value)
        for identity, values in simultaneous.items():
          memory[identity] = next(iter(values)) if len(values) == 1 else None
        if any(len(values) > 1 for values in simultaneous.values()):
          for event in self.events_by_onset[onset]:
            for tone in event.value['pitches']:
              if 'letter' in tone and len(simultaneous.get((tone['letter'], tone['octave']), ())) > 1:
                result[tone['id']] = True
    return result

def fragment_written(written: dict, original_start: fractions.Fraction, original_end: fractions.Fraction, start: fractions.Fraction, end: fractions.Fraction) -> dict:
  return written if (original_start, original_end) == (start, end) else generator.model.written_for(end - start)

def flags(written: dict) -> int:
  return max(0, generator.model.NOTE_TYPES.index(written['type']) - generator.model.NOTE_TYPES.index('quarter'))

def head_glyph(written: dict, head: dict | str | None) -> str:
  head = {'name': head} if isinstance(head, str) else head or {'name': 'normal'}
  if head.get('smufl'):
    return head['smufl']
  name = head['name']
  if name == 'none':
    return ''
  kind = written['type']
  suffix = 'DoubleWhole' if kind in ('maxima', 'long', 'breve') else 'Whole' if kind == 'whole' else 'Half' if kind == 'half' else 'Black'
  if head.get('filled') in ('yes', 'no'):
    suffix = 'Black' if head['filled'] == 'yes' else 'Half'
  if name == 'slash':
    return {'DoubleWhole': 'noteheadSlashWhiteDoubleWhole', 'Whole': 'noteheadSlashWhiteWhole', 'Half': 'noteheadSlashWhiteHalf', 'Black': 'noteheadSlashVerticalEnds'}[suffix]
  if name in ('slashed', 'back slashed'):
    return 'noteheadSlashed' + suffix + ('1' if name == 'slashed' else '2')
  if name in ('circled', 'circle dot'):
    return 'noteheadCircled' + suffix
  if name not in HEADS:
    raise generator.engraving.types.EngravingError(f'no SMuFL mapping for notehead {name!r}; specify its smufl glyph name')
  shape = HEADS[name]
  if shape in ('Moon', 'Round', 'Square', 'TriangleRight', 'TriangleLeft', 'TriangleRoundDown'):
    suffix = 'Black' if suffix == 'Black' else 'White'
  return 'notehead' + shape + suffix

def accidental(fonts: generator.engraving.fonts.Fonts, value: fractions.Fraction, display: dict, size: float, sources: tuple[str, ...]) -> list[generator.engraving.types.Primitive]:
  glyph = display.get('smufl') or ACCIDENTAL_GLYPHS.get(value)
  special = {'natural-sharp': ('accidentalNatural', 'accidentalSharp'), 'natural-flat': ('accidentalNatural', 'accidentalFlat')}
  if display.get('name') in special and not display.get('smufl'):
    items = []
    x = 0
    for name in special[display['name']]:
      part = fonts.glyph(name, x, 0, size, 'accidental', sources)
      items.append(part)
      x = part.box.right + 2
  elif glyph:
    items = [fonts.glyph(glyph, 0, 0, size, 'accidental', sources)]
  else:
    sign = '+' if value > 0 else ''
    items = [fonts.text(f'({sign}{value} st)', 0, 0, size * 1.15, 'accidental', sources)]
  if display.get('parentheses') == 'yes' or display.get('bracket') == 'yes':
    box = generator.engraving.types.bounds(items)
    left, right = ('[', ']') if display.get('bracket') == 'yes' else ('(', ')')
    items.extend((fonts.text(left, box.left - 7, size / 2, size * 1.8, 'accidental', sources), fonts.text(right, box.right + 2, size / 2, size * 1.8, 'accidental', sources)))
  return items