pipu-music/services/recommender/pipu_rec/analysis/features.py

"""Аудиопризнаки трека на numpy: тембр (эмбеддинг), темп, громкость, энергия, тональность.

Версия `basic-1` — без нейросетевых моделей (ADR-014): работает на любом CPU без TensorFlow.
Нейросетевые признаки Essentia (жанры, настроения, valence/arousal) — следующая версия анализатора;
смена `ANALYZER_VER` ставит пересчёт.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

ANALYZER_VER = "basic-1"
SAMPLE_RATE = 22050
N_FFT = 2048
HOP = 512
N_MELS = 64
EMBEDDING_DIM = 256
SEGMENT_S = 30.0

Float = NDArray[np.float32]


@dataclass(frozen=True)
class Features:
    bpm: float | None
    bpm_confidence: float
    key: str | None
    scale: str | None
    loudness_lufs: float
    energy: float
    danceability: float
    valence: float
    acousticness: float
    intro_ms: int
    embedding: Float  # (256,), L2-нормирован


def segments(samples: Float, sr: int = SAMPLE_RATE, seconds: float = SEGMENT_S) -> list[Float]:
    """Три фрагмента по 30 с — начало (после первых 10%), середина, конец; короткий трек — целиком."""
    n = int(seconds * sr)
    if len(samples) <= 3 * n:
        return [samples]
    total = len(samples)
    starts = [int(total * 0.1), total // 2 - n // 2, total - n - int(total * 0.05)]
    return [samples[s : s + n] for s in starts]


def _stft_mag(x: Float) -> Float:
    if len(x) < N_FFT:
        x = np.pad(x, (0, N_FFT - len(x)))
    frames = 1 + (len(x) - N_FFT) // HOP
    idx = np.arange(N_FFT)[None, :] + HOP * np.arange(frames)[:, None]
    window = np.hanning(N_FFT).astype(np.float32)
    return np.abs(np.fft.rfft(x[idx] * window, axis=1)).astype(np.float32)  # (frames, bins)


def _mel_filterbank(sr: int = SAMPLE_RATE) -> Float:
    def hz_to_mel(f: NDArray[np.float64]) -> NDArray[np.float64]:
        return 2595.0 * np.log10(1.0 + f / 700.0)

    def mel_to_hz(m: NDArray[np.float64]) -> NDArray[np.float64]:
        return 700.0 * (10 ** (m / 2595.0) - 1.0)

    bins = N_FFT // 2 + 1
    freqs = np.linspace(0, sr / 2, bins)
    mel_points = mel_to_hz(np.linspace(hz_to_mel(np.array(30.0)), hz_to_mel(np.array(sr / 2)), N_MELS + 2))
    fb = np.zeros((N_MELS, bins), dtype=np.float32)
    for m in range(N_MELS):
        lo, mid, hi = mel_points[m], mel_points[m + 1], mel_points[m + 2]
        up = (freqs - lo) / (mid - lo)
        down = (hi - freqs) / (hi - mid)
        fb[m] = np.maximum(0, np.minimum(up, down))
    return fb


_MEL = _mel_filterbank()

# Профили тональностей Крумханзла.
_MAJOR = 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])
_MINOR = 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])
_NOTES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def _log_mel(mag: Float) -> Float:
    """log-mel в дБ с полом на −60 дБ от максимума: шум в пустых полосах не превращается в «атаки»."""
    mel = mag**2 @ _MEL.T
    return (10 * np.log10(mel + 1e-6 * float(mel.max()) + 1e-12)).astype(np.float32)


def estimate_tempo(onset: Float, sr: int = SAMPLE_RATE) -> tuple[float | None, float]:
    """Темп по автокорреляции огибающей онсетов, 60–200 BPM с мягким приором к 120.
    Возвращает (bpm, уверенность 0..1); без выраженного ритма — (None, 0)."""
    fps = sr / HOP
    env = onset - onset.mean()
    # Ровный звук без атак: огибающая — численный шум, «ритма» в нём нет.
    if len(env) < int(fps * 4) or float(onset.std()) < 0.5:
        return None, 0.0
    ac = np.correlate(env, env, mode="full")[len(env) - 1 :]
    ac /= ac[0] + 1e-9
    lags = np.arange(len(ac))
    bpm_of = 60.0 * fps / np.maximum(lags, 1)
    ok = (bpm_of >= 60) & (bpm_of <= 200)
    if not np.any(ok):
        return None, 0.0
    prior = np.exp(-0.5 * (np.log2(bpm_of / 120.0) / 0.9) ** 2)
    weighted = np.where(ok, ac * prior, -np.inf)
    lag = int(np.argmax(weighted))
    # Уточнение по параболе вокруг пика.
    if 1 <= lag < len(ac) - 1:
        a, b, c = ac[lag - 1], ac[lag], ac[lag + 1]
        denom = a - 2 * b + c
        shift = 0.5 * (a - c) / denom if denom != 0 else 0.0
    else:
        shift = 0.0
    bpm = 60.0 * fps / (lag + shift)
    # Ясность пульса: насколько пик выше типичного уровня автокорреляции в диапазоне темпов.
    base = float(np.median(ac[ok]))
    confidence = float(np.clip((ac[lag] - base) / (1 - base + 1e-9), 0, 1))
    if confidence < 0.1:
        return None, confidence
    return round(float(bpm), 1), confidence


def _key(mag: Float, sr: int = SAMPLE_RATE) -> tuple[str | None, str | None, float]:
    """Тональность по хромаграмме; возвращает (нота, mode, уверенность)."""
    freqs = np.linspace(0, sr / 2, mag.shape[1])
    valid = (freqs > 60) & (freqs < 5000)
    pitch = np.round(12 * np.log2(freqs[valid] / 440.0) + 69).astype(int) % 12
    energy = mag[:, valid].sum(axis=0)
    chroma = np.bincount(pitch, weights=energy, minlength=12)
    if chroma.sum() <= 0:
        return None, None, 0.0
    chroma = chroma / chroma.sum()
    best = (-2.0, 0, "major")
    for shift in range(12):
        for name, prof in (("major", _MAJOR), ("minor", _MINOR)):
            r = float(np.corrcoef(chroma, np.roll(prof, shift))[0, 1])
            if r > best[0]:
                best = (r, shift, name)
    r, shift, scale = best
    return _NOTES[shift], scale, max(r, 0.0)


def analyze(samples: Float, sr: int = SAMPLE_RATE) -> Features:
    """Признаки трека по моно-сигналу sr Гц (float32, −1..1)."""
    if sr != SAMPLE_RATE:
        raise ValueError(f"expected {SAMPLE_RATE} Hz")
    parts = segments(samples, sr)
    mags = [_stft_mag(p) for p in parts]
    mag = np.concatenate(mags, axis=0)
    logmel = _log_mel(mag)  # (frames, mels)

    # Огибающая онсетов — положительный прирост log-mel; темп по каждому фрагменту, берём самый уверенный.
    tempos = []
    for m in mags:
        lm = _log_mel(m)
        onset = np.maximum(np.diff(lm, axis=0), 0).sum(axis=1).astype(np.float32)
        tempos.append(estimate_tempo(onset, sr))
    bpm, bpm_conf = max(tempos, key=lambda t: t[1])
    flux = np.maximum(np.diff(logmel, axis=0), 0).mean() if len(logmel) > 1 else 0.0

    rms = np.sqrt(np.mean(samples.astype(np.float64) ** 2) + 1e-12)
    loudness = float(-0.691 + 20 * np.log10(rms))  # без K-фильтра: близко к LUFS для музыки

    centroid = float((mag * np.arange(mag.shape[1])).sum() / (mag.sum() + 1e-9) / mag.shape[1])
    brightness = float(np.clip(centroid * 4, 0, 1))
    loud_part = np.clip((loudness + 30) / 22, 0, 1)
    energy = float(np.clip(0.5 * loud_part + 0.3 * np.clip(flux / 3, 0, 1) + 0.2 * brightness, 0, 1))
    danceability = float(np.clip(bpm_conf * 1.6, 0, 1)) if bpm else 0.0
    key, scale, key_conf = _key(mag, sr)
    mode = 1.0 if scale == "major" else -1.0 if scale == "minor" else 0.0
    valence = float(np.clip(0.45 + 0.2 * mode * key_conf + 0.35 * (brightness - 0.4), 0, 1))
    # Акустичность: мало высоких частот и плавная динамика.
    acousticness = float(np.clip(1 - brightness * 1.4 - flux / 6, 0, 1))

    # Вступление: время, пока громкость первого фрагмента не дойдёт до половины медианной.
    first = np.sqrt((mags[0] ** 2).mean(axis=1))
    above = np.nonzero(first >= 0.5 * np.median(first))[0]
    intro_ms = int(above[0] * HOP / sr * 1000) if len(above) else 0

    # Эмбеддинг: форма спектра (среднее без общей громкости), разброс, динамика по полосам, ритмичность.
    mean = logmel.mean(axis=0)
    shape = mean - mean.mean()
    std = logmel.std(axis=0)
    diff = np.diff(logmel, axis=0) if len(logmel) > 1 else np.zeros((1, N_MELS), np.float32)
    delta = np.abs(diff).mean(axis=0)
    band_flux = np.maximum(diff, 0).std(axis=0)
    emb = np.concatenate([shape / 10, std / 10, delta / 3, band_flux / 3]).astype(np.float32)
    emb = emb / (np.linalg.norm(emb) + 1e-9)

    return Features(
        bpm=bpm,
        bpm_confidence=round(bpm_conf, 3),
        key=key,
        scale=scale,
        loudness_lufs=round(loudness, 2),
        energy=round(energy, 3),
        danceability=round(danceability, 3),
        valence=round(valence, 3),
        acousticness=round(acousticness, 3),
        intro_ms=intro_ms,
        embedding=emb,
    )