product-vamp/generator/transposition.py

from __future__ import annotations

import copy
from dataclasses import dataclass
from typing import Any


@dataclass(frozen=True)
class KeyProfile:
  name: str
  root: int
  degrees: tuple[str, str, str, str, str, str, str]
  prefer_flats: bool


KEY_PROFILES = (
  KeyProfile("C", 0, ("C", "D", "E", "F", "G", "A", "B"), False),
  KeyProfile("Db", 1, ("Db", "Eb", "F", "Gb", "Ab", "Bb", "C"), True),
  KeyProfile("D", 2, ("D", "E", "F#", "G", "A", "B", "C#"), False),
  KeyProfile("Eb", 3, ("Eb", "F", "G", "Ab", "Bb", "C", "D"), True),
  KeyProfile("E", 4, ("E", "F#", "G#", "A", "B", "C#", "D#"), False),
  KeyProfile("F", 5, ("F", "G", "A", "Bb", "C", "D", "E"), True),
  KeyProfile("Gb", 6, ("Gb", "Ab", "Bb", "Cb", "Db", "Eb", "F"), True),
  KeyProfile("G", 7, ("G", "A", "B", "C", "D", "E", "F#"), False),
  KeyProfile("Ab", 8, ("Ab", "Bb", "C", "Db", "Eb", "F", "G"), True),
  KeyProfile("A", 9, ("A", "B", "C#", "D", "E", "F#", "G#"), False),
  KeyProfile("Bb", 10, ("Bb", "C", "D", "Eb", "F", "G", "A"), True),
  KeyProfile("B", 11, ("B", "C#", "D#", "E", "F#", "G#", "A#"), False),
)
TARGET_KEYS = tuple(profile.name for profile in KEY_PROFILES)
NATURAL_PITCH_CLASSES = {"c": 0, "d": 2, "e": 4, "f": 5, "g": 7, "a": 9, "b": 11}
LETTER_DEGREES = {"c": 0, "d": 1, "e": 2, "f": 3, "g": 4, "a": 5, "b": 6}


def mod12(value: int) -> int:
  return value % 12


def find_key(name: str) -> KeyProfile:
  normalized = name.lower()
  aliases = {"c#": "db", "d#": "eb", "f#": "gb", "g#": "ab", "a#": "bb"}
  normalized = aliases.get(normalized, normalized)
  for profile in KEY_PROFILES:
    if profile.name.lower() == normalized:
      return profile
  raise ValueError(f"unsupported key signature: {name}")


def signed_interval(source: KeyProfile, target: KeyProfile) -> int:
  interval = mod12(target.root - source.root)
  return interval - 12 if interval > 6 else interval


def split_spelling(value: str) -> tuple[str, str]:
  if not value:
    return "C", ""
  return value[0], value[1:]


def accidental_offset(accidental: str) -> int:
  value = accidental.lower()
  if value in {"#", "sharp"}:
    return 1
  if value in {"b", "flat"}:
    return -1
  return 0


def spelling_pitch_class(spelling: str) -> int:
  letter, accidental = split_spelling(spelling)
  return mod12(NATURAL_PITCH_CLASSES.get(letter.lower(), 0) + accidental_offset(accidental))


def pitch_class(pitch: dict[str, Any]) -> int:
  letter = str(pitch.get("letter", "c")).lower()
  accidental = str(pitch.get("accidental", "n"))
  return mod12(NATURAL_PITCH_CLASSES.get(letter, 0) + accidental_offset(accidental))


def pitch_midi(pitch: dict[str, Any]) -> int:
  octave = int(pitch.get("octave", 4))
  return octave * 12 + pitch_class(pitch)


def octave_for(midi: int, pitch_class_value: int) -> int:
  octave = (midi - pitch_class_value) // 12
  while octave * 12 + pitch_class_value > midi:
    octave -= 1
  while (octave + 1) * 12 + pitch_class_value <= midi:
    octave += 1
  return octave


def matching_source_degree_for_pitch(pitch: dict[str, Any], source: KeyProfile) -> int | None:
  actual = pitch_class(pitch)
  source_letter_degree = LETTER_DEGREES.get(str(pitch.get("letter", "c")).lower(), 0)
  for index, spelling in enumerate(source.degrees):
    letter, _ = split_spelling(spelling)
    if LETTER_DEGREES.get(letter.lower(), 0) == source_letter_degree and spelling_pitch_class(spelling) == actual:
      return index
  return None


def chromatic_spelling(pitch_class_value: int, prefer_flats: bool) -> str:
  flats = ("C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B")
  sharps = ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B")
  return flats[pitch_class_value] if prefer_flats else sharps[pitch_class_value]


def transpose_pitch(pitch: dict[str, Any], source: KeyProfile, target: KeyProfile, interval: int) -> dict[str, Any]:
  transposed_midi = pitch_midi(pitch) + interval
  degree = matching_source_degree_for_pitch(pitch, source)
  spelling = target.degrees[degree] if degree is not None else chromatic_spelling(mod12(transposed_midi), target.prefer_flats)
  letter, accidental = split_spelling(spelling)
  pc = spelling_pitch_class(spelling)
  return {
    "letter": letter.lower(),
    "accidental": accidental or "n",
    "octave": octave_for(transposed_midi, pc),
  }


def parse_chord_root(value: str, offset: int) -> tuple[int, str] | None:
  if offset >= len(value) or value[offset].upper() not in "ABCDEFG":
    return None
  length = 1
  if offset + length < len(value) and value[offset + length] in {"#", "b"}:
    length += 1
  return length, value[offset:offset + length]


def matching_source_degree_for_root(root: str, source: KeyProfile) -> int | None:
  root_letter, _ = split_spelling(root)
  actual = spelling_pitch_class(root)
  for index, spelling in enumerate(source.degrees):
    letter, _ = split_spelling(spelling)
    if LETTER_DEGREES.get(letter.lower(), 0) == LETTER_DEGREES.get(root_letter.lower(), 0) and spelling_pitch_class(spelling) == actual:
      return index
  return None


def transpose_root(root: str, source: KeyProfile, target: KeyProfile, interval: int) -> str:
  degree = matching_source_degree_for_root(root, source)
  if degree is not None:
    return target.degrees[degree]
  return chromatic_spelling(mod12(spelling_pitch_class(root) + interval), target.prefer_flats)


def transpose_chord_symbol(symbol: str, source: KeyProfile, target: KeyProfile, interval: int) -> str:
  root = parse_chord_root(symbol, 0)
  if root is None:
    return symbol
  root_length, root_text = root
  result = transpose_root(root_text, source, target, interval) + symbol[root_length:]
  slash = result.find("/")
  if slash < 0:
    return result
  bass = parse_chord_root(result, slash + 1)
  if bass is None:
    return result
  bass_length, bass_text = bass
  return result[:slash + 1] + transpose_root(bass_text, source, target, interval) + result[slash + 1 + bass_length:]


def transpose_document(document: dict[str, Any], target_key_name: str) -> dict[str, Any]:
  output = copy.deepcopy(document)
  metadata = output.setdefault("metadata", {})
  if not isinstance(metadata, dict):
    raise ValueError("score metadata must be an object")
  source_key = find_key(str(metadata.get("key_signature", "C") or "C"))
  target_key = find_key(target_key_name)
  interval = signed_interval(source_key, target_key)
  metadata["key_signature"] = target_key.name
  for staff in output.get("staves", []):
    if not isinstance(staff, dict):
      continue
    for measure in staff.get("measures", []):
      if not isinstance(measure, dict):
        continue
      measure["key_signature"] = target_key.name
      for voice in measure.get("voices", []):
        if not isinstance(voice, dict):
          continue
        for note in voice.get("notes", []):
          if not isinstance(note, dict):
            continue
          chord = note.get("chord_symbol")
          if isinstance(chord, str) and chord:
            note["chord_symbol"] = transpose_chord_symbol(chord, source_key, target_key, interval)
          pitches = note.get("pitches")
          if isinstance(pitches, list):
            note["pitches"] = [transpose_pitch(pitch, source_key, target_key, interval) if isinstance(pitch, dict) else pitch for pitch in pitches]
  return output