theme_generation/common.py

"""Common symbols and musical constants for generation engines."""

from __future__ import annotations

from dataclasses import dataclass


PITCH_CLASS = {
    "C": 0,
    "B#": 0,
    "C#": 1,
    "Db": 1,
    "D": 2,
    "D#": 3,
    "Eb": 3,
    "E": 4,
    "Fb": 4,
    "E#": 5,
    "F": 5,
    "F#": 6,
    "Gb": 6,
    "G": 7,
    "G#": 8,
    "Ab": 8,
    "A": 9,
    "A#": 10,
    "Bb": 10,
    "B": 11,
    "Cb": 11,
}

PC_TO_SHARP_NAME = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

COMMON_DURATION_BEATS = {
    "64th": 0.0625,
    "32nd": 0.125,
    "16th triplet": 1 / 6,
    "16th": 0.25,
    "eighth triplet": 1 / 3,
    "dotted 16th": 0.375,
    "eighth": 0.5,
    "quarter triplet": 2 / 3,
    "dotted eighth": 0.75,
    "quarter": 1.0,
    "dotted quarter": 1.5,
    "half": 2.0,
    "dotted half": 3.0,
    "whole": 4.0,
}

ABC_DURATION = {
    "64th": "/4",
    "32nd": "/2",
    "16th triplet": "/3",
    "16th": "",
    "eighth triplet": "4/3",
    "dotted 16th": "3/2",
    "eighth": "2",
    "quarter triplet": "8/3",
    "dotted eighth": "3",
    "quarter": "4",
    "dotted quarter": "6",
    "half": "8",
    "dotted half": "12",
    "whole": "16",
}

DEFAULT_DURATIONS = {
    "16th",
    "eighth",
    "quarter",
    "half",
    "dotted eighth",
    "dotted quarter",
    "eighth triplet",
    "16th triplet",
}

START_SYMBOL = "__THEME_GENERATION_START__"


@dataclass(frozen=True)
class Symbol:
    """A key-relative pitch class plus notated duration."""

    rpc: int
    duration: str

    def __str__(self) -> str:
        return f"{self.rpc}:{self.duration}"


def duration_options(min_duration: str, duration_grid: str, allow_triplets: bool) -> set[str]:
    if min_duration not in COMMON_DURATION_BEATS:
        labels = ", ".join(sorted(COMMON_DURATION_BEATS))
        raise ValueError(f"Unknown duration {min_duration!r}; choose one of: {labels}")
    if duration_grid not in COMMON_DURATION_BEATS:
        labels = ", ".join(sorted(COMMON_DURATION_BEATS))
        raise ValueError(f"Unknown duration grid {duration_grid!r}; choose one of: {labels}")

    min_beats = COMMON_DURATION_BEATS[min_duration]
    grid_beats = COMMON_DURATION_BEATS[duration_grid]
    allowed = set()
    for duration in DEFAULT_DURATIONS:
        duration_beats = COMMON_DURATION_BEATS[duration]
        on_grid = abs((duration_beats / grid_beats) - round(duration_beats / grid_beats)) < 1e-9
        regular_triplet = allow_triplets and "triplet" in duration
        if duration_beats >= min_beats and (on_grid or regular_triplet):
            allowed.add(duration)
    return allowed


def duration_units(duration: str) -> int:
    return round(COMMON_DURATION_BEATS[duration] * 12)


def is_triplet_duration(duration: str) -> bool:
    return "triplet" in duration


def parse_key_root(raw_key: str | None) -> int | None:
    if not raw_key:
        return None
    root = raw_key.split("%", 1)[0].strip().replace("m", "")
    return PITCH_CLASS.get(root)