product-vamp/generator/engraving/layout.py

from __future__ import annotations

import bisect
import collections
import dataclasses
import fractions
import math

import generator.engraving.attributes
import generator.engraving.audit
import generator.engraving.fonts
import generator.engraving.notation
import generator.engraving.notes
import generator.engraving.scales
import generator.engraving.spans
import generator.engraving.tablature
import generator.engraving.types
import generator.harmony
import generator.model

def lyric_text(verse: dict) -> str:
  return ''.join(s['text'] for s in verse['syllables'] if s['kind'] in ('text', 'elision'))

def column(index: generator.engraving.notation.ScoreIndex, time: fractions.Fraction, start: fractions.Fraction, end: fractions.Fraction, fonts: generator.engraving.fonts.Fonts, style: generator.engraving.types.Style) -> tuple[list[generator.engraving.types.Primitive], list[generator.engraving.types.Anchor], list[generator.engraving.types.Fragment]]:
  items, anchors, fragments = [], [], []
  events = index.events_at(time, include_sustained=time == start)
  grace_width = 0
  occupied = []
  for event in events:
    try:
      drawn, placed, pieces = generator.engraving.notes.draw_event(index, event, start, end, fonts, style)
    except ValueError as error:
      raise generator.engraving.types.EngravingError(f'measure {event.measure["number"]}, voice {event.voice}, {event.value["id"]}: {error}') from error
    x = grace_width
    collidable = [p for p in drawn if p.role in ('head', 'tab', 'rest', 'slash', 'accidental', 'dot')]
    while any(p.box.moved(x, 0).overlaps(q.box, 2) for p in collidable for q in occupied):
      x += style.space * 1.4
    drawn = [p.moved(x, 0) for p in drawn]
    placed = [a.moved(x, 0) for a in placed]
    if event.value.get('grace'):
      grace_width = max((p.box.right for p in drawn), default=grace_width) + 9
    occupied.extend(p for p in drawn if p.role in ('head', 'tab', 'rest', 'slash', 'accidental', 'dot'))
    items.extend(drawn)
    items.extend(generator.engraving.tablature.column(index, event, placed, fonts, style))
    anchors.extend(placed)
    fragments.extend(pieces)
    if event.start == time and not event.value.get('hidden'):
      verses = event.value.get('lyrics', [])
      for verse in verses:
        lane = index.lyric_lanes.index((event.voice, verse['number']))
        text = lyric_text(verse)
        if text:
          items.append(fonts.text(text, x + style.space / 2, lane * 23, 18, 'lyric', (event.value['id'] + ':lyric:' + verse['number'], event.value['id']), True))
  measure_start, _, measure = index.measure_at(time)
  for owner, _, item in index.harmonies_by_onset.get(time, []):
    if owner is measure:
      items.append(fonts.text(generator.harmony.symbol(item), 0, 0, 18 if index.staff.get('tablature') else 21, 'harmony', (item['id'],)))
  items.extend(generator.engraving.scales.column(index, time, start, fonts))
  if time == measure_start:
    items.append(fonts.text(str(measure['number']), 0, -style.space * (index.state(time)['staff_details']['lines'] + 0.5), 12, 'measure_number', (measure['id'],)))
    if label := measure.get('section_label'):
      items.append(fonts.text(label, 0, 0, 16 if index.staff.get('tablature') else 19, 'section', (measure['id'] + ':section',)))
  return items, anchors, fragments

def column_metrics(index, time, start, end, fonts, style, columns):
  events = index.events_at(time, include_sustained=time == start)
  finish = min(end, max((event.end for event in events), default=time))
  key = (time, start, finish)
  if columns is not None and key in columns:
    columns.move_to_end(key)
    return columns[key]
  drawn, placed, pieces = column(index, time, start, end, fonts, style)
  result = drawn, placed, pieces, generator.engraving.types.bounds(drawn)
  if columns is not None:
    columns[key] = result
    if len(columns) > 256:
      columns.popitem(last=False)
  return result

def make_system(index: generator.engraving.notation.ScoreIndex, start: fractions.Fraction, end: fractions.Fraction, fonts: generator.engraving.fonts.Fonts, style: generator.engraving.types.Style, justify: bool = False, *, columns=None, staff_width: float | None = None) -> generator.engraving.types.System:
  items, anchors, fragments = [], [], []
  state = index.state(start)
  measure = index.measure_at(start)[2]
  header = generator.engraving.attributes.staff_attributes(index, state, start, fonts, style.space, measure['id'] + ':initial', ('clef', 'key', 'meter', 'staff_details'))
  items.extend(header)
  cursor = max((p.box.right for p in header), default=0) + 22
  content_left = cursor - 9
  times = [start] + index.column_times[bisect.bisect_right(index.column_times, start):bisect.bisect_left(index.column_times, end)]
  placements = []
  groups = []
  for ti, time in enumerate(times):
    item_start, anchor_start = len(items), len(anchors)
    drawn, placed, pieces, box = column_metrics(index, time, start, end, fonts, style, columns)
    leading = []
    measure_start, _, m = index.measure_at(time)
    if time == measure_start and time > start:
      previous = index.measure_at(time, before=True)[2]
      leading.extend(generator.engraving.attributes.barline(previous.get('barlines', {}).get('right', 'regular'), 0, -style.space * (index.state(time)['staff_details']['lines'] - 1), previous['id']))
      old_state = index.state(time, before=True)
      fields = tuple(k for k in ('clef', 'key', 'meter', 'staff_details') if old_state[k] != index.state(time)[k])
      leading.extend(p.moved(12, 0) for p in generator.engraving.attributes.staff_attributes(index, index.state(time), time, fonts, style.space, m['id'] + ':initial', fields))
    if time > start:
      for owner, ci, change in index.attributes_by_onset.get(time, []):
        extra = generator.engraving.attributes.staff_attributes(index, index.state(time), time, fonts, style.space, owner['id'] + f':attributes:{ci}', tuple(change))
        leading.extend(p.moved(max((p.box.right for p in leading), default=0) + 10, 0) for p in extra)
    if leading:
      width = generator.engraving.types.bounds(leading).right - generator.engraving.types.bounds(leading).left + 14
      items.extend(p.moved(cursor - generator.engraving.types.bounds(leading).left, 0) for p in leading)
      cursor += width
    x = cursor - min(0, box.left)
    placements.append((time, x))
    items.extend(p.moved(x, 0) for p in drawn)
    anchors.extend(a.moved(x, 0) for a in placed)
    fragments.extend(pieces)
    next_time = times[ti + 1] if ti + 1 < len(times) else end
    spring = style.column_gap + 13 * math.sqrt(min(4, float(next_time - time)))
    cursor = x + max(style.space * 1.4, box.right) + spring
    groups.append((item_start, len(items), anchor_start, len(anchors)))
  width = cursor + 6
  available = staff_width + 1 if staff_width is not None else style.width - style.margin_x * 2
  if justify and width < available:
    stretch = (available - width) / max(1, len(times))
    for i, (first_item, last_item, first_anchor, last_anchor) in enumerate(groups):
      items[first_item:last_item] = [p.moved(i * stretch, 0) for p in items[first_item:last_item]]
      anchors[first_anchor:last_anchor] = [a.moved(i * stretch, 0) for a in anchors[first_anchor:last_anchor]]
    placements = [(time, x + i * stretch) for i, (time, x) in enumerate(placements)]
    width = available
  if staff_width is not None:
    if width > available + 0.01:
      raise generator.engraving.types.EngravingError(f'at musical offset {start}: content exceeds the final staff width')
    width = available
  right = width - 5
  last_measure = index.measure_at(end, before=True)[2]
  tab_items = [p for p in items if p.role.startswith('tab')]
  items = [p for p in items if not p.role.startswith('tab')]
  items = generator.engraving.spans.draw_beams(index, anchors, items, start, end, content_left, right, style)
  items = generator.engraving.spans.draw_spans(index, anchors, items, start, end, content_left, right, fonts, style)
  items = generator.engraving.spans.place_text_lanes(items, fonts, style)
  items = generator.engraving.spans.draw_lyric_spans(index, items, anchors, start, end, content_left, right, fonts)
  changes = [(0, state['staff_details']['lines'])]
  for time, x in placements:
    lines = index.state(time)['staff_details']['lines']
    if changes[-1][1] != lines:
      changes.append((x - 8, lines))
  staff_lines = []
  for i, (x1, lines) in enumerate(changes):
    x2 = changes[i + 1][0] if i + 1 < len(changes) else width - 1
    for number in range(lines):
      y = -number * style.space
      staff_lines.append(generator.engraving.types.line('staff', (), ((x1, y), (x2, y)), style.space * 0.13))
  items = staff_lines + items
  if end in index.boundaries:
    items.extend(generator.engraving.attributes.barline(last_measure.get('barlines', {}).get('right', 'regular'), width - 1, -style.space * (index.state(end, before=True)['staff_details']['lines'] - 1), last_measure['id']))
  if start in index.boundaries:
    items.extend(generator.engraving.attributes.barline(measure.get('barlines', {}).get('left', 'none'), 0, -style.space * (state['staff_details']['lines'] - 1), measure['id']))
  else:
    items.append(fonts.text(str(measure['number']) + ' (cont.)', 0, min(p.box.top for p in items) - 6, 12, 'continuation', (measure['id'],)))
  items = generator.engraving.tablature.staff(index, items + tab_items, start, end, content_left, width - 1, fonts, style)
  return generator.engraving.types.System(start, end, tuple(items), tuple(anchors), tuple(fragments), generator.engraving.types.bounds(items))

def fit_part_systems(systems, index, available_height, fonts, style):

  def fitting_scale(items):
    scale = 1
    for first in range(0, len(items), 2):
      pair = items[first:first + 2]
      height = sum(s.box.bottom - s.box.top for s in pair)
      scale = min(scale, (available_height - style.system_gap * (len(pair) - 1)) / height)
    return scale

  scale = fitting_scale(systems)
  columns = collections.OrderedDict()
  # Fit the entire part together; only spacing, never page-specific glyph sizes, varies.
  for _ in range(8):
    width = (style.width - 2 * style.margin_x - 2) / scale
    fitted = [make_system(index, s.start, s.end, fonts, style, justify=s.end < index.end, columns=columns, staff_width=width) for s in systems]
    next_scale = fitting_scale(fitted)
    if scale <= next_scale + 0.000001:
      return [s.scaled(scale) for s in fitted]
    scale = next_scale
  raise generator.engraving.types.EngravingError('part spacing did not converge')

def part_header(document, fonts, style):
  size = 16
  title = document['metadata']['title']
  item = fonts.text(title, style.margin_x, style.margin_y + size, size, 'title', ())
  while item.box.right > style.width - style.margin_x and size > 12:
    size -= 1
    item = fonts.text(title, style.margin_x, style.margin_y + size, size, 'title', ())
  if item.box.right > style.width - style.margin_x:
    raise generator.engraving.types.EngravingError('score title cannot fit the page without clipping')
  result = [item]
  subtitle = ' - '.join(str(value) for value in (document['metadata'].get('composer'), document['staves'][0].get('name')) if value)
  if not subtitle:
    return result
  credit = fonts.text(subtitle, 0, style.margin_y + size, 12, 'subtitle', ())
  width = credit.box.right - credit.box.left
  if item.box.right + 20 + width <= style.width - style.margin_x:
    result.append(credit.moved(style.width - style.margin_x - credit.box.right, 0))
    return result
  line = ''
  baseline = item.box.bottom + 15
  for word in subtitle.split():
    text = ' '.join(filter(None, (line, word)))
    candidate = fonts.text(text, style.margin_x, baseline, 12, 'subtitle', ())
    if candidate.box.right > style.width - style.margin_x and line:
      result.append(fonts.text(line, style.margin_x, baseline, 12, 'subtitle', ()))
      baseline = result[-1].box.bottom + 15
      text = word
      candidate = fonts.text(text, style.margin_x, baseline, 12, 'subtitle', ())
    if candidate.box.right > style.width - style.margin_x:
      raise generator.engraving.types.EngravingError('score credits cannot fit the page without clipping')
    line = text
  result.append(fonts.text(line, style.margin_x, baseline, 12, 'subtitle', ()))
  return result

def paginate(systems: list[generator.engraving.types.System], document: dict, fonts: generator.engraving.fonts.Fonts, style: generator.engraving.types.Style, *, index=None) -> tuple[generator.engraving.types.Page, ...]:
  pages = []
  index = index or generator.engraving.notation.ScoreIndex(document)
  header = part_header(document, fonts, style)
  top = generator.engraving.types.bounds(header).bottom + style.page_clearance
  footers = [fonts.text(str(number), style.width / 2, style.height - style.margin_y, 12, 'page_number', (), True) for number in range(1, (len(systems) + 1) // 2 + 1)]
  bottom = min(p.box.top for p in footers) - style.page_clearance
  if bottom - top <= style.system_gap:
    raise generator.engraving.types.EngravingError('page header leaves no room for two music systems')
  fitted = fit_part_systems(systems, index, bottom - top, fonts, style)
  for first in range(0, len(systems), 2):
    number = len(pages) + 1
    pair = fitted[first:first + 2]
    slack = max(0, bottom - top - sum(s.box.bottom - s.box.top for s in pair) - style.system_gap * (len(pair) - 1))
    spacing = slack / 3 if len(pair) == 2 else 0
    cursor = top + spacing
    placed = []
    for system in pair:
      dx, dy = style.margin_x + 1, cursor - system.box.top
      placed.append(dataclasses.replace(system, primitives=tuple(p.moved(dx, dy) for p in system.primitives), anchors=tuple(a.moved(dx, dy) for a in system.anchors), box=system.box.moved(dx, dy)))
      cursor = placed[-1].box.bottom + style.system_gap + spacing
    footer = footers[first // 2]
    pages.append(generator.engraving.types.Page(number, style.width, style.height, style.supersampling, tuple(placed), tuple(header + [p for s in placed for p in s.primitives] + [footer])))
  return tuple(pages)

def layout_systems(document: dict, style: generator.engraving.types.Style, fonts: generator.engraving.fonts.Fonts, *, index=None) -> tuple[generator.engraving.types.System, ...]:
  if index is None:
    generator.model.validate_document(document)
    index = generator.engraving.notation.ScoreIndex(document)
  if not index.measures:
    raise generator.engraving.types.EngravingError('score has no measures')
  systems = []
  start = fractions.Fraction()
  available = style.width - style.margin_x * 2 - 2
  while start < index.end:
    columns = collections.OrderedDict()
    candidates = index.boundary_times[bisect.bisect_right(index.boundary_times, start):]
    selected = None
    # Reserve title/credit and footer bands before fitting complete measures.
    height_limit = style.height - style.margin_y * 2 - 110
    for end in candidates:
      trial = make_system(index, start, end, fonts, style, columns=columns)
      if trial.box.right - trial.box.left > available or trial.box.bottom - trial.box.top > height_limit:
        break
      selected = trial
    if selected is None:
      measure = index.measure_at(start)[2]
      raise generator.engraving.types.EngravingError(f'measure {measure["number"]} ({measure["id"]}): complete measure exceeds the page width or height at readable spacing')
    systems.append(selected)
    start = selected.end
  return tuple(systems)

def layout_score(document: dict, style: generator.engraving.types.Style, fonts: generator.engraving.fonts.Fonts) -> generator.engraving.types.LayoutDocument:
  generator.model.validate_document(document)
  style = generator.engraving.types.score_style(document, style)
  index = generator.engraving.notation.ScoreIndex(document)
  systems = layout_systems(document, style, fonts, index=index)
  layout = generator.engraving.types.LayoutDocument(paginate(systems, document, fonts, style, index=index), style, fonts.fingerprint(style))
  generator.engraving.audit.audit_layout(document, layout, fonts)
  return layout