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 document


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"


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)