core/part_splitter.py

"""
core/part_splitter.py

NoteEvent 리스트를 N개 파트로 분배.

현재 전략: Greedy 피치 범위 기반 분배
  - Part 1: 가장 높은 피치 그룹 (주 멜로디)
  - Part 2: 중간 피치 그룹 (화음/반주)
  - Part 3: 가장 낮은 피치 그룹 (베이스)

교체 가능성: 이 함수의 시그니처를 유지하면서
  더 정교한 알고리즘(staff/voice 기반, 음역대 분석 등)으로 교체 가능.
"""

from __future__ import annotations

from .models import NoteEvent


def _count_max_simultaneous(notes: list[NoteEvent]) -> int:
    """최대 동시 발음 수 계산 (part_count 자동 감지용)."""
    if not notes:
        return 1
    events: list[tuple[float, int]] = []
    for n in notes:
        events.append((n.start, 1))
        events.append((n.start + n.duration, -1))
    events.sort()
    max_count = current = 0
    for _, delta in events:
        current += delta
        max_count = max(max_count, current)
    return max(1, max_count)


def split_parts(
    notes: list[NoteEvent],
    part_count: int = 0,
) -> list[list[NoteEvent]]:
    """
    NoteEvent 리스트를 part_count개 파트로 분배.

    Args:
        notes: 전체 NoteEvent 리스트
        part_count: 분배할 파트 수 (기본 3)

    Returns:
        파트별 NoteEvent 리스트의 리스트.
        길이는 항상 part_count이며, 비어있는 파트는 빈 리스트.
    """
    if part_count <= 0:
        part_count = _count_max_simultaneous(notes)

    if not notes:
        return [[] for _ in range(part_count)]

    # part_hint가 지정된 경우 우선 사용
    if any(n.part_hint is not None for n in notes):
        return _split_by_hint(notes, part_count)

    staffs = {n.staff for n in notes}
    voices = {n.voice for n in notes}

    # staff도 1개이고 voice도 1개면 단선율 → Part 1에 전부
    if len(staffs) == 1 and len(voices) == 1:
        parts = [[] for _ in range(part_count)]
        parts[0].extend(sorted(notes, key=lambda n: n.start))
        return parts

    # staff 또는 voice가 여러 개면 staff+voice 기반 분배
    return _split_by_staff_and_voice(notes, part_count)

    # fallback: 피치 범위 기반 greedy 분배
    return _split_by_pitch_range(notes, part_count)


def _split_by_hint(notes: list[NoteEvent], part_count: int) -> list[list[NoteEvent]]:
    """part_hint 값에 따라 분배."""
    parts: list[list[NoteEvent]] = [[] for _ in range(part_count)]
    for note in notes:
        hint = note.part_hint
        if hint is not None and 1 <= hint <= part_count:
            parts[hint - 1].append(note)
        else:
            parts[0].append(note)  # 미지정은 Part 1으로
    return parts


def _split_by_staff_and_voice(notes: list[NoteEvent], part_count: int) -> list[list[NoteEvent]]:
    """
    staff + voice 조합으로 파트 분배.

    조합 순서 (높은 피치 우선): staff1/voice1, staff1/voice2, staff2/voice1, ...
    """
    # (staff, voice) 조합을 평균 피치 기준으로 정렬 (높은 것이 Part 1)
    combos: dict[tuple[int, int], list[NoteEvent]] = {}
    for note in notes:
        key = (note.staff, note.voice)
        combos.setdefault(key, []).append(note)

    def avg_pitch(note_list: list[NoteEvent]) -> float:
        pitched = [n.pitch for n in note_list if n.pitch > 0]
        return sum(pitched) / len(pitched) if pitched else 0.0

    sorted_combos = sorted(combos.values(), key=avg_pitch, reverse=True)

    parts: list[list[NoteEvent]] = [[] for _ in range(part_count)]
    for i, group in enumerate(sorted_combos):
        idx = min(i, part_count - 1)
        parts[idx].extend(group)

    # 파트 내 start 기준 재정렬
    for part in parts:
        part.sort(key=lambda n: n.start)

    return parts


def _split_by_pitch_range(notes: list[NoteEvent], part_count: int) -> list[list[NoteEvent]]:
    """
    피치 범위 기준 greedy 분배.

    전체 피치를 part_count 구간으로 나누어 분배.
    """
    pitched = [n for n in notes if n.pitch > 0]
    rests = [n for n in notes if n.pitch == 0]

    if not pitched:
        parts: list[list[NoteEvent]] = [[] for _ in range(part_count)]
        parts[part_count - 1].extend(rests)
        return parts

    min_pitch = min(n.pitch for n in pitched)
    max_pitch = max(n.pitch for n in pitched)
    pitch_span = max_pitch - min_pitch or 1

    parts = [[] for _ in range(part_count)]
    for note in pitched:
        ratio = (note.pitch - min_pitch) / pitch_span
        # 높은 피치 = 낮은 인덱스 (Part 1)
        idx = part_count - 1 - int(ratio * (part_count - 0.01))
        idx = max(0, min(idx, part_count - 1))
        parts[idx].append(note)

    # 쉼표는 Part 1에 배치
    parts[0].extend(rests)

    for part in parts:
        part.sort(key=lambda n: n.start)

    return parts