product-vamp/generator/score.py

from __future__ import annotations

import copy
import json
import re
from pathlib import Path
from typing import Any

WHOLE_TICKS = 16384
HALF_TICKS = WHOLE_TICKS // 2
QUARTER_TICKS = WHOLE_TICKS // 4
EIGHTH_TICKS = WHOLE_TICKS // 8
SIXTEENTH_TICKS = WHOLE_TICKS // 16
THIRTY_SECOND_TICKS = WHOLE_TICKS // 32
RENDERABLE_TICKS = (
  WHOLE_TICKS + HALF_TICKS,
  WHOLE_TICKS,
  HALF_TICKS + QUARTER_TICKS,
  HALF_TICKS,
  QUARTER_TICKS + EIGHTH_TICKS,
  QUARTER_TICKS,
  EIGHTH_TICKS + SIXTEENTH_TICKS,
  EIGHTH_TICKS,
  SIXTEENTH_TICKS + THIRTY_SECOND_TICKS,
  SIXTEENTH_TICKS,
  THIRTY_SECOND_TICKS,
)

REST_PITCH = {"letter": "b", "accidental": "n", "octave": 4}
FORBIDDEN_OUTPUT_FIELDS = {
  "tab_positions",
  "tab_frets",
  "finger",
  "fingers",
  "fingering",
  "left_hand_finger",
  "right_hand_finger",
}


def load_json(path: Path) -> dict[str, Any]:
  with path.open("r", encoding="utf-8") as stream:
    document = json.load(stream)
  if not isinstance(document, dict):
    raise ValueError(f"{path} is not a score document")
  return normalize_score_document(document, path)


def write_json(path: Path, document: dict[str, Any]) -> None:
  path.parent.mkdir(parents=True, exist_ok=True)
  with path.open("w", encoding="utf-8") as stream:
    json.dump(document, stream, indent=2, ensure_ascii=False)
    stream.write("\n")


def slugify(value: str) -> str:
  value = value.strip().lower()
  value = re.sub(r"[^a-z0-9]+", "-", value)
  value = value.strip("-")
  return value or "score"


def human_label(value: str) -> str:
  text = value.replace("_", " ").replace("-", " ").strip()
  return " ".join(part.capitalize() for part in text.split()) or "General"


class CompactScoreError(ValueError):
  pass


DURATION_TICKS = {
  "1": WHOLE_TICKS,
  "2": HALF_TICKS,
  "4": QUARTER_TICKS,
  "8": EIGHTH_TICKS,
  "16": SIXTEENTH_TICKS,
  "32": THIRTY_SECOND_TICKS,
}


def score_field(document: dict[str, Any], metadata: dict[str, Any], *names: str, fallback: str = "") -> str:
  for name in names:
    value = document.get(name)
    if isinstance(value, str):
      return value
    value = metadata.get(name)
    if isinstance(value, str):
      return value
  return fallback


def parse_duration_token(text: str, label: str) -> int:
  dotted = text.endswith(".")
  base_text = text[:-1] if dotted else text
  if base_text not in DURATION_TICKS:
    raise CompactScoreError(f"{label}: unsupported duration {text!r}")
  ticks = DURATION_TICKS[base_text]
  return ticks + ticks // 2 if dotted else ticks


def compact_pitch(letter: str, accidental: str, octave: int) -> dict[str, Any]:
  return {
    "letter": letter.lower(),
    "accidental": accidental or "n",
    "octave": octave,
  }


class CompactMeasureParser:
  def __init__(self, text: str, *, measure_number: int, time_signature: str) -> None:
    self.text = text
    self.index = 0
    self.measure_number = measure_number
    self.time_signature = time_signature
    self.sequence = 0
    self.tuplet_sequence = 0
    self.pending_chord = ""
    self.last_note: dict[str, Any] | None = None
    self.last_octave = 4
    self.tuplets: list[dict[str, Any]] = []

  def label(self) -> str:
    return f"measure {self.measure_number}"

  def error(self, message: str) -> CompactScoreError:
    return CompactScoreError(f"{self.label()}, column {self.index + 1}: {message}")

  def peek(self) -> str:
    return self.text[self.index] if self.index < len(self.text) else ""

  def skip_space(self) -> None:
    while self.peek() and self.peek().isspace():
      self.index += 1

  def parse_until(self, close: str, label: str) -> str:
    start = self.index
    escaped = False
    while self.index < len(self.text):
      char = self.text[self.index]
      if escaped:
        escaped = False
        self.index += 1
        continue
      if char == "\\":
        escaped = True
        self.index += 1
        continue
      if char == close:
        value = self.text[start:self.index]
        self.index += 1
        return bytes(value, "utf-8").decode("unicode_escape")
      self.index += 1
    raise self.error(f"unterminated {label}")

  def parse_chord(self) -> None:
    self.index += 1
    self.pending_chord = self.parse_until("]", "chord symbol").strip()
    if not self.pending_chord:
      raise self.error("chord symbol must not be empty")

  def parse_lyric(self) -> None:
    if self.last_note is None:
      raise self.error("lyric must follow a note")
    self.index += 1
    lyric = self.parse_until('"', "lyric")
    self.last_note["lyric"] = lyric

  def parse_tuplet_header(self) -> tuple[int, int]:
    self.skip_space()
    match = re.match(r"(\d+)\s*:\s*(\d+)", self.text[self.index:])
    if match is None:
      raise self.error("tuplet must start with N:M")
    self.index += match.end()
    note_count = int(match.group(1))
    space_count = int(match.group(2))
    if note_count <= 0 or space_count <= 0:
      raise self.error("tuplet counts must be positive")
    return note_count, space_count

  def parse_tuplet(self) -> list[dict[str, Any]]:
    self.index += 1
    note_count, space_count = self.parse_tuplet_header()
    self.tuplet_sequence += 1
    tuplet_id = f"tuplet-{self.measure_number}-{self.tuplet_sequence}"
    notes: list[dict[str, Any]] = []
    while True:
      self.skip_space()
      if not self.peek():
        raise self.error("unterminated tuplet")
      if self.peek() == "}":
        self.index += 1
        break
      if self.peek() == "{":
        raise self.error("nested tuplets are not supported")
      notes.extend(self.parse_item())
    if len(notes) != note_count:
      raise self.error(f"tuplet {tuplet_id} contains {len(notes)} notes/rests, expected {note_count}")
    self.apply_tuplet_duration(notes, note_count, space_count)
    for note in notes:
      note["tuplet_id"] = tuplet_id
    self.tuplets.append({
      "id": tuplet_id,
      "note_count": note_count,
      "space_count": space_count,
    })
    return notes

  def apply_tuplet_duration(self, notes: list[dict[str, Any]], note_count: int, space_count: int) -> None:
    numerators = [int(note["duration_ticks"]) * space_count for note in notes]
    floors = [value // note_count for value in numerators]
    target_total = sum(numerators) // note_count
    remainder = target_total - sum(floors)
    order = sorted(range(len(notes)), key=lambda index: numerators[index] % note_count, reverse=True)
    actual = floors[:]
    for index in order[:remainder]:
      actual[index] += 1
    for note, ticks in zip(notes, actual):
      display = int(note["duration_ticks"])
      note["duration_ticks"] = max(1, ticks)
      note["display_ticks"] = display

  def parse_pitch_atom(self) -> dict[str, Any]:
    start = self.index
    char = self.peek()
    if char.lower() not in "abcdefg":
      raise self.error("expected pitch")
    self.index += 1
    accidental = ""
    if self.peek() in {"#", "b", "n"}:
      accidental = self.peek()
      self.index += 1
    octave_start = self.index
    while self.peek().isdigit():
      self.index += 1
    if self.index > octave_start:
      self.last_octave = int(self.text[octave_start:self.index])
    if self.index == start:
      raise self.error("expected pitch")
    return compact_pitch(char, accidental, self.last_octave)

  def parse_pitch_group(self) -> list[dict[str, Any]]:
    self.index += 1
    pitches: list[dict[str, Any]] = []
    while True:
      self.skip_space()
      if not self.peek():
        raise self.error("unterminated pitch group")
      if self.peek() == ")":
        self.index += 1
        break
      pitches.append(self.parse_pitch_atom())
    if not pitches:
      raise self.error("pitch group must contain at least one pitch")
    return pitches

  def parse_note(self, tie_from_previous: bool = False) -> dict[str, Any]:
    start = self.index
    char = self.peek()
    if not char:
      raise self.error("expected note or rest")
    kind = "rest" if char.lower() == "r" else "note"
    if kind == "rest":
      self.index += 1
      pitches = [copy.deepcopy(REST_PITCH)]
    elif char == "(":
      pitches = self.parse_pitch_group()
    else:
      pitches = [self.parse_pitch_atom()]
    if self.peek() != ":":
      raise self.error(f"expected ':' after {self.text[start:self.index]!r}")
    self.index += 1
    duration_start = self.index
    while self.peek().isdigit():
      self.index += 1
    if self.peek() == ".":
      self.index += 1
    duration_text = self.text[duration_start:self.index]
    if not duration_text:
      raise self.error("duration is required")
    self.sequence += 1
    note: dict[str, Any] = {
      "id": f"note-{self.measure_number}-{self.sequence}",
      "kind": kind,
      "duration_ticks": parse_duration_token(duration_text, self.label()),
      "pitches": pitches,
      "hidden": False,
    }
    if self.pending_chord:
      note["chord_symbol"] = self.pending_chord
      self.pending_chord = ""
    if tie_from_previous:
      note["tie_from_previous"] = True
    if self.peek() == "~":
      note["tie_to_next"] = True
      self.index += 1
    self.last_note = note
    return note

  def parse_item(self) -> list[dict[str, Any]]:
    self.skip_space()
    char = self.peek()
    if not char:
      return []
    if char == "[":
      self.parse_chord()
      return []
    if char == '"':
      self.parse_lyric()
      return []
    if char == "{":
      return self.parse_tuplet()
    if char == "}":
      return []
    if char == "~":
      self.index += 1
      self.skip_space()
      next_char = self.peek()
      if next_char and (next_char.lower() in "abcdefgr" or next_char == "("):
        return [self.parse_note(tie_from_previous=True)]
      if self.last_note is None:
        raise self.error("tie marker must attach to a note")
      self.last_note["tie_to_next"] = True
      return []
    if char.lower() in "abcdefgr" or char == "(":
      return [self.parse_note()]
    raise self.error("expected note, rest, chord, lyric, tie, or tuplet")

  def parse(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    notes: list[dict[str, Any]] = []
    while True:
      self.skip_space()
      if not self.peek():
        break
      notes.extend(self.parse_item())
    if self.pending_chord:
      raise self.error(f"chord {self.pending_chord!r} is not followed by a note or rest")
    total = sum(int(note.get("duration_ticks", 0)) for note in notes)
    capacity = measure_capacity_ticks(self.time_signature)
    if total > capacity:
      raise CompactScoreError(f"{self.label()}: duration total {total} exceeds {self.time_signature} capacity {capacity}")
    return notes, self.tuplets


def compact_measure_text(value: Any, index: int) -> tuple[str, dict[str, str]]:
  if isinstance(value, str):
    return value, {}
  if isinstance(value, dict):
    text = value.get("notes", value.get("notation", ""))
    if not isinstance(text, str):
      raise CompactScoreError(f"measure {index + 1}: notes must be a string")
    overrides: dict[str, str] = {}
    for source, target in (
      ("section", "section_label"),
      ("section_label", "section_label"),
      ("key", "key_signature"),
      ("key_signature", "key_signature"),
      ("time", "time_signature"),
      ("time_signature", "time_signature"),
    ):
      field = value.get(source)
      if isinstance(field, str):
        overrides[target] = field
    return text, overrides
  raise CompactScoreError(f"measure {index + 1}: compact measure must be a string or object")


def normalize_score_document(document: dict[str, Any], path: Path | None = None) -> dict[str, Any]:
  if isinstance(document.get("staves"), list):
    return document
  measures_value = document.get("measures")
  if not isinstance(measures_value, list):
    return document
  metadata_value = document.get("metadata", {})
  metadata = metadata_value if isinstance(metadata_value, dict) else {}
  title = score_field(document, metadata, "title", fallback=human_label(path.stem if path is not None else "score"))
  composer = score_field(document, metadata, "composer")
  key_signature = score_field(document, metadata, "key_signature", "key", fallback="C")
  time_signature = score_field(document, metadata, "time_signature", "time", fallback="4/4")
  clef = score_field(document, metadata, "clef", fallback="treble")
  part = score_field(document, metadata, "part", "name", fallback="Vocal")
  role = score_field(document, metadata, "role", fallback="vocal")

  output_metadata = sanitize_json(metadata)
  if not isinstance(output_metadata, dict):
    output_metadata = {}
  output_metadata.update({
    "title": title,
    "composer": composer,
    "key_signature": key_signature,
    "time_signature": time_signature,
  })
  for key, value in document.items():
    if key in {"title", "composer", "key", "key_signature", "time", "time_signature", "clef", "part", "name", "role", "measures", "metadata"}:
      continue
    output_metadata[key] = sanitize_json(value)

  measures: list[dict[str, Any]] = []
  for index, value in enumerate(measures_value):
    text, overrides = compact_measure_text(value, index)
    measure_time = overrides.get("time_signature", time_signature)
    measure_key = overrides.get("key_signature", key_signature)
    parser = CompactMeasureParser(text, measure_number=index + 1, time_signature=measure_time)
    notes, tuplets = parser.parse()
    measures.append({
      "id": f"measure-{index + 1}",
      "section_label": overrides.get("section_label", ""),
      "time_signature": measure_time,
      "key_signature": measure_key,
      "voices": [{
        "id": f"voice-{index + 1}",
        "notes": notes,
        "tuplets": tuplets,
      }],
    })

  return {
    "version": int(document.get("version", 1)) if isinstance(document.get("version", 1), int) else 1,
    "metadata": output_metadata,
    "staves": [{
      "id": "staff-source",
      "name": part,
      "clef": clef,
      "role": role,
      "measures": measures,
    }],
  }


def source_staff(document: dict[str, Any]) -> dict[str, Any]:
  staves = document.get("staves")
  if not isinstance(staves, list) or not staves:
    raise ValueError("score document must contain at least one staff")
  first_staff: dict[str, Any] | None = None
  for value in staves:
    if not isinstance(value, dict):
      continue
    if first_staff is None:
      first_staff = value
    staff_id = str(value.get("id", ""))
    if not value.get("arranged") and staff_id not in {"staff-arranged-guitar", "staff-arranged-vocal", "staff-arranged-practice"} and value.get("role") != "guitar":
      return value
  if first_staff is None:
    raise ValueError("score document does not contain an object staff")
  return first_staff


def measure_capacity_ticks(time_signature: str) -> int:
  try:
    beats_text, unit_text = time_signature.split("/", 1)
    beats = int(beats_text)
    unit = int(unit_text)
  except ValueError:
    return WHOLE_TICKS
  if beats <= 0 or unit <= 0:
    return WHOLE_TICKS
  return WHOLE_TICKS * beats // unit


def note_duration(note: dict[str, Any]) -> int:
  value = note.get("duration_ticks", QUARTER_TICKS)
  return max(1, int(value)) if isinstance(value, int) else QUARTER_TICKS


def renderable_duration(duration_ticks: int) -> bool:
  return duration_ticks in RENDERABLE_TICKS


def renderable_duration_parts(duration_ticks: int) -> list[int]:
  parts: list[int] = []
  remaining = duration_ticks
  while remaining > 0:
    match = next((duration for duration in RENDERABLE_TICKS if duration <= remaining), None)
    if match is None:
      return [duration_ticks]
    parts.append(match)
    remaining -= match
  return parts


def first_voice(measure: dict[str, Any]) -> dict[str, Any]:
  voices = measure.get("voices")
  if isinstance(voices, list) and voices and isinstance(voices[0], dict):
    return voices[0]
  return {"id": "voice-empty", "tuplets": [], "notes": []}


def first_chord_symbol(measure: dict[str, Any], fallback: str = "") -> str:
  voice = first_voice(measure)
  notes = voice.get("notes")
  if not isinstance(notes, list):
    return fallback
  for note in notes:
    if isinstance(note, dict):
      chord = note.get("chord_symbol")
      if isinstance(chord, str) and chord:
        return chord
  return fallback


def measure_notes(measure: dict[str, Any]) -> list[dict[str, Any]]:
  voice = first_voice(measure)
  notes = voice.get("notes")
  return [note for note in notes if isinstance(note, dict)] if isinstance(notes, list) else []


def timed_measure_notes(measure: dict[str, Any]) -> list[tuple[int, dict[str, Any]]]:
  cursor = 0
  result: list[tuple[int, dict[str, Any]]] = []
  for note in measure_notes(measure):
    result.append((cursor, note))
    cursor += note_duration(note)
  return result


def chord_markers(measure: dict[str, Any]) -> list[tuple[int, int, str]]:
  markers: list[tuple[int, int, str]] = []
  for start_ticks, note in timed_measure_notes(measure):
    chord = note.get("chord_symbol")
    if isinstance(chord, str) and chord:
      markers.append((start_ticks, note_duration(note), chord))
  return markers


def chord_symbol_at_offset(measure: dict[str, Any], offset_ticks: int) -> str:
  for start_ticks, _, chord in chord_markers(measure):
    if start_ticks == offset_ticks:
      return chord
  return ""


def active_chord_at_offset(measure: dict[str, Any], offset_ticks: int, fallback: str = "") -> str:
  current = fallback
  for start_ticks, _, chord in chord_markers(measure):
    if start_ticks > offset_ticks:
      break
    current = chord
  return current


def source_note_attack_at(measure: dict[str, Any], offset_ticks: int) -> bool:
  for start_ticks, note in timed_measure_notes(measure):
    if start_ticks == offset_ticks and note.get("kind", "note") == "note" and isinstance(note.get("pitches"), list) and note["pitches"]:
      return True
  return False


def first_pitched_note_start(measure: dict[str, Any]) -> int | None:
  cursor = 0
  for note in measure_notes(measure):
    if note.get("kind", "note") == "note" and isinstance(note.get("pitches"), list) and note["pitches"]:
      return cursor
    cursor += note_duration(note)
  return None


def leading_pickup_measure_index(measures: list[dict[str, Any]]) -> int | None:
  if len(measures) < 2 or first_chord_symbol(measures[0]):
    return None
  first_pitch = first_pitched_note_start(measures[0])
  if first_pitch is None or first_pitch <= 0:
    return None
  for index, measure in enumerate(measures):
    if first_chord_symbol(measure):
      return 0 if index == 1 else None
  return None


def body_start_measure_index(measures: list[dict[str, Any]]) -> int:
  pickup = leading_pickup_measure_index(measures)
  return pickup + 1 if pickup is not None else 0


def rest_note(note_id: str, duration_ticks: int, chord_symbol: str = "") -> dict[str, Any]:
  note: dict[str, Any] = {
    "id": note_id,
    "kind": "rest",
    "duration_ticks": duration_ticks,
    "pitches": [copy.deepcopy(REST_PITCH)],
    "hidden": False,
  }
  if chord_symbol:
    note["chord_symbol"] = chord_symbol
  return note


def renderable_rest_notes(note_id_prefix: str, duration_ticks: int, sequence: int, chord_symbol: str = "") -> tuple[list[dict[str, Any]], int]:
  notes: list[dict[str, Any]] = []
  for index, duration in enumerate(renderable_duration_parts(duration_ticks)):
    sequence += 1
    notes.append(rest_note(f"{note_id_prefix}-{sequence}", duration, chord_symbol if index == 0 else ""))
  return notes, sequence


def apply_source_chord_markers(notes: list[dict[str, Any]], measure: dict[str, Any]) -> list[dict[str, Any]]:
  markers = [(start, chord) for start, _, chord in chord_markers(measure)]
  if not markers:
    result = [copy.deepcopy(note) for note in notes]
    for note in result:
      note.pop("chord_symbol", None)
    return result

  unique_markers: list[tuple[int, str]] = []
  for start, chord in markers:
    if unique_markers and unique_markers[-1][0] == start:
      unique_markers[-1] = (start, chord)
    else:
      unique_markers.append((start, chord))

  marker_by_start = dict(unique_markers)
  marker_starts = [start for start, _ in unique_markers]
  result: list[dict[str, Any]] = []
  cursor = 0
  split_sequence = 0
  for source_note in notes:
    note = copy.deepcopy(source_note)
    note.pop("chord_symbol", None)
    duration = note_duration(note)
    end = cursor + duration
    split_points = [cursor] + [start for start in marker_starts if cursor < start < end] + [end]
    for index in range(len(split_points) - 1):
      start = split_points[index]
      stop = split_points[index + 1]
      part_duration = stop - start
      if part_duration <= 0:
        continue
      part = copy.deepcopy(note)
      part["duration_ticks"] = part_duration
      if len(split_points) > 2:
        split_sequence += 1
        part["id"] = f"{note.get('id', 'note')}-split-{split_sequence}"
        if part.get("kind", "note") == "note":
          part["tie_to_next"] = index + 2 < len(split_points)
          part["tie_from_previous"] = index > 0
      if start in marker_by_start:
        part["chord_symbol"] = marker_by_start[start]
      result.append(part)
    cursor = end
  return result


def slash_note(note_id: str, duration_ticks: int, chord_symbol: str = "") -> dict[str, Any]:
  note: dict[str, Any] = {
    "id": note_id,
    "kind": "note",
    "duration_ticks": duration_ticks,
    "pitches": [copy.deepcopy(REST_PITCH)],
    "notehead": "slash",
    "hidden": False,
  }
  if chord_symbol:
    note["chord_symbol"] = chord_symbol
  return note


def sanitize_json(value: Any) -> Any:
  if isinstance(value, dict):
    return {key: sanitize_json(item) for key, item in value.items() if key not in FORBIDDEN_OUTPUT_FIELDS}
  if isinstance(value, list):
    return [sanitize_json(item) for item in value]
  return copy.deepcopy(value)


def sanitized_note(note: dict[str, Any], *, strip_lyrics: bool) -> dict[str, Any]:
  output = sanitize_json(note)
  if strip_lyrics:
    output.pop("lyric", None)
  output.setdefault("hidden", False)
  return output


def copy_source_voice(measure: dict[str, Any], *, strip_lyrics: bool) -> dict[str, Any]:
  voice = first_voice(measure)
  tuplets = sanitize_json(voice.get("tuplets", []))
  notes = [sanitized_note(note, strip_lyrics=strip_lyrics) for note in measure_notes(measure)]
  return {
    "id": str(voice.get("id", "voice-source")),
    "tuplets": tuplets if isinstance(tuplets, list) else [],
    "notes": notes or [rest_note("arr-rest-empty", measure_capacity_ticks(str(measure.get("time_signature", "4/4"))))],
  }


def make_measure(measure_id: str, source: dict[str, Any], notes: list[dict[str, Any]], *, section_label: str = "", tuplets: list[Any] | None = None) -> dict[str, Any]:
  measure: dict[str, Any] = {
    "id": measure_id,
    "time_signature": str(source.get("time_signature", "4/4")),
    "key_signature": str(source.get("key_signature", "C")),
    "voices": [{
      "id": f"voice-{measure_id}",
      "tuplets": tuplets or [],
      "notes": notes,
    }],
  }
  if section_label:
    measure["section_label"] = section_label
  return measure


def make_arranged_document(source_document: dict[str, Any], staff: dict[str, Any], *, staff_id: str, staff_name: str, role: str, arrangement_part_type: str, measures: list[dict[str, Any]]) -> dict[str, Any]:
  metadata = sanitize_json(source_document.get("metadata", {}))
  if not isinstance(metadata, dict):
    metadata = {}
  output: dict[str, Any] = {
    "version": int(source_document.get("version", 1)) if isinstance(source_document.get("version", 1), int) else 1,
    "metadata": metadata,
    "staves": [{
      "id": staff_id,
      "name": staff_name,
      "clef": str(staff.get("clef", "treble")),
      "role": role,
      "arranged": True,
      "arrangement_part_type": arrangement_part_type,
      "measures": measures,
    }],
  }
  return sanitize_json(output)


def count_json_files(root: Path) -> int:
  return sum(1 for path in root.rglob("*.json") if path.is_file())


def validate_generated_tree(root: Path) -> None:
  for path in root.rglob("*.json"):
    text = path.read_text(encoding="utf-8")
    for field in FORBIDDEN_OUTPUT_FIELDS:
      if f'"{field}"' in text:
        raise ValueError(f"{path} contains forbidden generated field {field}")
    json.loads(text)