python-basic-pitch/tonal_inference.py

from __future__ import annotations

from dataclasses import dataclass
import os
import time
from typing import Any, Optional

import librosa
import numpy as np
from scipy import signal as scipy_signal

NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
MAJOR_PROFILE = np.array(
    [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88],
    dtype=np.float32,
)
MINOR_PROFILE = np.array(
    [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17],
    dtype=np.float32,
)
FIFTHS_ORDER = ["C", "G", "D", "A", "E", "B", "F#", "C#", "G#", "D#", "A#", "F"]
FIFTHS_INDEX = {name: index for index, name in enumerate(FIFTHS_ORDER)}
MAX_PITCH_EVIDENCE_SECONDS = 12.0
PITCH_EVIDENCE_TARGET_SR = 11025


@dataclass
class TonalPitchEvidence:
    histogram: np.ndarray
    dominant_pc: Optional[int]
    dominant_ratio: float
    onset_histogram: np.ndarray


@dataclass
class TonalValidationCandidate:
    tonic: str
    mode: str
    score: float
    confidence: float
    acoustic_correlation: float
    root_support: float
    tonic_pitch_support: float
    dominant_pitch_support: float
    diatonic_coverage: float
    circle_coherence: float


def extract_tonal_pitch_evidence(
    audio: np.ndarray,
    sr: int,
    faixa: tuple[float, float],
    hop_length: int = 512,
    base_histogram: Optional[np.ndarray] = None,
    onset_histogram_hint: Optional[np.ndarray] = None,
    prefer_fast_mode: bool = False,
) -> TonalPitchEvidence:
    stage_timing_enabled = os.getenv("AUDIO_REPORT_STAGE_TIMING", "").strip().lower() in {"1", "true", "yes", "on"}

    def mark_stage(label: str, started_at: float) -> float:
        if not stage_timing_enabled:
            return time.perf_counter()
        now = time.perf_counter()
        print(f"[tonal_timing] {label} ({now - started_at:.2f}s)", flush=True)
        return now

    data = np.asarray(audio, dtype=np.float32)
    if data.size == 0:
      return TonalPitchEvidence(
          histogram=np.zeros(12, dtype=np.float32),
          dominant_pc=None,
          dominant_ratio=0.0,
          onset_histogram=np.zeros(12, dtype=np.float32),
      )

    if prefer_fast_mode and base_histogram is not None:
        return build_histogram_pitch_evidence(
            base_histogram=base_histogram,
            onset_histogram_hint=onset_histogram_hint,
        )

    data, sr = compact_pitch_evidence_audio(data, sr)
    stage_started_at = mark_stage("compact_audio", time.perf_counter())

    n_fft = 2048
    try:
        stft = np.abs(
            librosa.stft(
                y=data,
                n_fft=n_fft,
                hop_length=hop_length,
                center=True,
            ),
        ).astype(np.float32)
    except Exception:
        return TonalPitchEvidence(
            histogram=np.zeros(12, dtype=np.float32),
            dominant_pc=None,
            dominant_ratio=0.0,
            onset_histogram=np.zeros(12, dtype=np.float32),
        )
    stage_started_at = mark_stage("spectral_tracking", stage_started_at)

    freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft).astype(np.float32)
    mask = (
        (freqs >= max(40.0, float(faixa[0]) * 0.75))
        & (freqs <= min(float(sr) / 2.0 - 1.0, float(faixa[1]) * 1.35))
    )
    if not np.any(mask):
        return TonalPitchEvidence(
            histogram=np.zeros(12, dtype=np.float32),
            dominant_pc=None,
            dominant_ratio=0.0,
            onset_histogram=np.zeros(12, dtype=np.float32),
        )

    masked_freqs = freqs[mask]
    masked_stft = stft[mask, :]

    histogram = np.zeros(12, dtype=np.float32)
    onset_histogram = np.zeros(12, dtype=np.float32)
    n_frames = masked_stft.shape[1] if masked_stft.ndim == 2 else 0

    for frame_index in range(n_frames):
        frame_magnitudes = masked_stft[:, frame_index]
        if frame_magnitudes.size == 0:
            continue

        peak = float(np.max(frame_magnitudes))
        if peak <= 1e-8:
            continue

        top_indices = np.argpartition(frame_magnitudes, -4)[-4:]
        frame_total = 0.0
        strongest_pc: Optional[int] = None
        strongest_weight = 0.0
        for idx in top_indices:
            freq = float(masked_freqs[idx])
            mag = float(frame_magnitudes[idx])
            if not np.isfinite(freq) or not np.isfinite(mag) or mag <= peak * 0.22:
                continue
            if freq < max(40.0, float(faixa[0]) * 0.72) or freq > min(float(sr) / 2.0 - 1.0, float(faixa[1]) * 1.4):
                continue
            midi = 69.0 + 12.0 * np.log2(freq / 440.0)
            pc = int(round(midi)) % 12
            clarity = max(0.0, mag / peak)
            weight = mag * (0.55 + clarity)
            histogram[pc] += weight
            frame_total += weight
            if weight > strongest_weight:
                strongest_weight = weight
                strongest_pc = pc

        if strongest_pc is not None and frame_total > 0:
            onset_histogram[strongest_pc] += strongest_weight / frame_total
    stage_started_at = mark_stage("histogram_accumulation", stage_started_at)

    histogram = normalize_vector(histogram)
    onset_histogram = normalize_vector(onset_histogram)
    if float(histogram.sum()) <= 0:
        return TonalPitchEvidence(
            histogram=histogram,
            dominant_pc=None,
            dominant_ratio=0.0,
            onset_histogram=onset_histogram,
        )

    dominant_pc = int(np.argmax(histogram))
    dominant_ratio = float(histogram[dominant_pc])
    return TonalPitchEvidence(
        histogram=histogram,
        dominant_pc=dominant_pc,
        dominant_ratio=dominant_ratio,
        onset_histogram=onset_histogram,
    )


def build_histogram_pitch_evidence(
    base_histogram: Optional[np.ndarray],
    onset_histogram_hint: Optional[np.ndarray] = None,
) -> TonalPitchEvidence:
    histogram = normalize_vector(np.asarray(base_histogram if base_histogram is not None else np.zeros(12), dtype=np.float32))
    onset_histogram = normalize_vector(
        np.asarray(onset_histogram_hint if onset_histogram_hint is not None else histogram, dtype=np.float32),
    )
    if float(histogram.sum()) <= 0:
        return TonalPitchEvidence(
            histogram=np.zeros(12, dtype=np.float32),
            dominant_pc=None,
            dominant_ratio=0.0,
            onset_histogram=np.zeros(12, dtype=np.float32),
        )

    dominant_pc = int(np.argmax(histogram))
    dominant_ratio = float(histogram[dominant_pc])
    return TonalPitchEvidence(
        histogram=histogram,
        dominant_pc=dominant_pc,
        dominant_ratio=dominant_ratio,
        onset_histogram=onset_histogram,
    )


def compact_pitch_evidence_audio(audio: np.ndarray, sr: int) -> tuple[np.ndarray, int]:
    data = np.asarray(audio, dtype=np.float32)
    if data.size == 0:
        return data, sr

    if sr > PITCH_EVIDENCE_TARGET_SR:
        data = resample_audio(data, sr, PITCH_EVIDENCE_TARGET_SR)
        sr = PITCH_EVIDENCE_TARGET_SR

    max_samples = int(MAX_PITCH_EVIDENCE_SECONDS * sr)
    if data.size <= max_samples:
        return data, sr

    frame = max(1024, int(sr * 0.5))
    hop = max(512, int(sr * 0.25))
    energies: list[tuple[float, int, int]] = []
    for start in range(0, max(1, data.size - frame), hop):
        end = min(data.size, start + frame)
        chunk = data[start:end]
        if chunk.size < frame // 2:
            continue
        rms = float(np.sqrt(np.mean(np.square(chunk)) + 1e-10))
        energies.append((rms, start, end))

    if not energies:
        return data[:max_samples], sr

    energies.sort(key=lambda item: item[0], reverse=True)
    selected: list[tuple[int, int]] = []
    target_windows = max(2, min(3, max_samples // max(frame, 1)))
    for _rms, start, end in energies:
        if any(abs(start - existing_start) < int(sr * 2.5) for existing_start, _existing_end in selected):
            continue
        selected.append((start, end))
        if len(selected) >= target_windows:
            break

    if not selected:
        return data[:max_samples], sr

    selected.sort(key=lambda item: item[0])
    snippets = [data[start:end] for start, end in selected]
    compacted = np.concatenate(snippets).astype(np.float32)
    if compacted.size > max_samples:
        compacted = compacted[:max_samples]
    return compacted, sr


def resample_audio(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    if orig_sr == target_sr or audio.size == 0:
        return np.asarray(audio, dtype=np.float32)
    gcd = np.gcd(int(orig_sr), int(target_sr))
    up = int(target_sr // gcd)
    down = int(orig_sr // gcd)
    return scipy_signal.resample_poly(audio, up, down).astype(np.float32)


def validate_key_candidates(
    base_candidates: list[Any],
    chroma_mean: np.ndarray,
    root_histogram: np.ndarray,
    acoustic_events: list[dict[str, Any]],
    pitch_evidence: TonalPitchEvidence,
    note_to_pitch_class,
    correlation_pearson,
    logistic,
) -> list[TonalValidationCandidate]:
    validated: list[TonalValidationCandidate] = []

    first_root = note_to_pitch_class(str(acoustic_events[0].get("nome", ""))) if acoustic_events else None
    last_root = note_to_pitch_class(str(acoustic_events[-1].get("nome", ""))) if acoustic_events else None

    for candidate in base_candidates:
        tonic = str(candidate.tonic)
        mode = str(candidate.mode)
        tonic_pc = NOTE_NAMES.index(tonic)
        profile = np.roll(MAJOR_PROFILE if mode == "maior" else MINOR_PROFILE, tonic_pc)
        acoustic_correlation = float(correlation_pearson(chroma_mean, profile))
        root_support = float(root_histogram[tonic_pc]) if root_histogram.size == 12 else 0.0
        dominant_pc = (tonic_pc + 7) % 12
        subdominant_pc = (tonic_pc + 5) % 12
        relative_pc = (tonic_pc + (9 if mode == "maior" else 3)) % 12

        dominant_root_support = float(root_histogram[dominant_pc]) if root_histogram.size == 12 else 0.0
        subdominant_root_support = float(root_histogram[subdominant_pc]) if root_histogram.size == 12 else 0.0
        tonic_pitch_support = float(pitch_evidence.histogram[tonic_pc]) if pitch_evidence.histogram.size == 12 else 0.0
        dominant_pitch_support = float(pitch_evidence.histogram[dominant_pc]) if pitch_evidence.histogram.size == 12 else 0.0
        relative_pitch_support = float(pitch_evidence.histogram[relative_pc]) if pitch_evidence.histogram.size == 12 else 0.0
        onset_tonic_support = float(pitch_evidence.onset_histogram[tonic_pc]) if pitch_evidence.onset_histogram.size == 12 else 0.0

        diatonic_coverage = weighted_diatonic_coverage(
            acoustic_events,
            tonic_pc=tonic_pc,
            mode=mode,
            note_to_pitch_class=note_to_pitch_class,
        )
        circle_coherence = weighted_circle_of_fifths_coherence(
            acoustic_events,
            tonic_pc=tonic_pc,
            note_to_pitch_class=note_to_pitch_class,
        )

        edge_bonus = 0.0
        if first_root == tonic_pc:
            edge_bonus += 0.3
        if last_root == tonic_pc:
            edge_bonus += 0.52
        elif last_root == dominant_pc:
            edge_bonus += 0.18

        dominant_hint = 0.0
        if pitch_evidence.dominant_pc is not None:
            if pitch_evidence.dominant_pc == tonic_pc:
                dominant_hint += 0.32
            elif pitch_evidence.dominant_pc == dominant_pc:
                dominant_hint += 0.22
            elif pitch_evidence.dominant_pc == relative_pc:
                dominant_hint += 0.08

        final_score = (
            float(candidate.score) * 0.18
            + acoustic_correlation * 3.25
            + root_support * 4.9
            + dominant_root_support * 2.15
            + subdominant_root_support * 0.95
            + tonic_pitch_support * 4.15
            + dominant_pitch_support * 1.9
            + onset_tonic_support * 1.35
            + relative_pitch_support * 0.4
            + diatonic_coverage * 2.55
            + circle_coherence * 1.45
            + edge_bonus
            + dominant_hint
        )
        confidence = float(logistic(final_score / 7.4))
        validated.append(
            TonalValidationCandidate(
                tonic=tonic,
                mode=mode,
                score=float(final_score),
                confidence=confidence,
                acoustic_correlation=acoustic_correlation,
                root_support=root_support,
                tonic_pitch_support=tonic_pitch_support,
                dominant_pitch_support=dominant_pitch_support,
                diatonic_coverage=diatonic_coverage,
                circle_coherence=circle_coherence,
            )
        )

    validated.sort(key=lambda item: item.score, reverse=True)
    return validated[:6]


def weighted_diatonic_coverage(
    events: list[dict[str, Any]],
    tonic_pc: int,
    mode: str,
    note_to_pitch_class,
) -> float:
    if not events:
        return 0.0

    diatonic = diatonic_pitch_classes(tonic_pc, mode)
    supported = 0.0
    total = 0.0
    for event in events:
        root = note_to_pitch_class(str(event.get("nome", "")))
        if root is None:
            continue
        weight = max(
            0.18,
            (float(event.get("fim", 0.0)) - float(event.get("inicio", 0.0)))
            * max(0.15, float(event.get("confianca", 0.0))),
        )
        total += weight
        if root in diatonic:
            supported += weight
    if total <= 0:
        return 0.0
    return supported / total


def weighted_circle_of_fifths_coherence(
    events: list[dict[str, Any]],
    tonic_pc: int,
    note_to_pitch_class,
) -> float:
    if not events:
        return 0.0

    anchors = {
        tonic_pc,
        (tonic_pc + 7) % 12,
        (tonic_pc + 5) % 12,
        (tonic_pc + 2) % 12,
        (tonic_pc + 10) % 12,
    }
    total = 0.0
    score = 0.0
    for event in events:
        root = note_to_pitch_class(str(event.get("nome", "")))
        if root is None:
            continue
        weight = max(
            0.18,
            (float(event.get("fim", 0.0)) - float(event.get("inicio", 0.0)))
            * max(0.15, float(event.get("confianca", 0.0))),
        )
        total += weight
        best_distance = min(circle_distance(root, anchor) for anchor in anchors)
        score += weight * max(0.0, 1.0 - (best_distance / 4.0))
    if total <= 0:
        return 0.0
    return score / total


def diatonic_pitch_classes(tonic_pc: int, mode: str) -> set[int]:
    intervals = [0, 2, 3, 5, 7, 8, 10] if mode == "menor" else [0, 2, 4, 5, 7, 9, 11]
    return {int((tonic_pc + interval) % 12) for interval in intervals}


def circle_distance(left_pc: int, right_pc: int) -> int:
    left_name = NOTE_NAMES[int(left_pc) % 12]
    right_name = NOTE_NAMES[int(right_pc) % 12]
    left_index = FIFTHS_INDEX[left_name]
    right_index = FIFTHS_INDEX[right_name]
    diff = abs(left_index - right_index)
    return min(diff, 12 - diff)


def normalize_vector(values: np.ndarray) -> np.ndarray:
    vector = np.asarray(values, dtype=np.float32)
    total = float(vector.sum())
    if total <= 1e-8:
        return np.zeros_like(vector)
    return vector / total