"""Акустические признаки реплики.

На записи с одним микрофоном на столе участники сидят на разном расстоянии,
и это слышно: дальний голос глуше и тише. Тембровые модели такую разницу
почти не улавливают, поэтому признаки считаются отдельно и отдаются наружу -
по ним LLM может уточнить, кто где, когда разделение по голосу ненадёжно.
"""
import numpy as np

__all__ = ["segment_pitch", "pitch_split", "segment_acoustics", "separation_quality"]

SAMPLE_RATE = 16000


# Диапазон человеческого голоса с запасом: женский обычно 165-255 Гц,
# мужской 85-155 Гц.
F0_MIN, F0_MAX = 70.0, 320.0
F0_FRAME_SEC = 0.04
F0_HOP_SEC = 0.02
# Ниже этой корреляции кадр считаем шумом, а не голосом.
F0_VOICED = 0.35
# Разрыв между группами, ниже которого голоса считаем неразличимыми.
# На восьми записях: смешанные пары дают 68-93 Гц, однополые 11-33 Гц.
PITCH_SPLIT_HZ = 40.0


def frame_pitch(frame: np.ndarray) -> float | None:
    """Основная частота кадра по автокорреляции.

    Голос периодичен, и период виден как пик автокорреляции. Пик ниже порога
    означает, что периодичности нет: шум, шелест, пауза.
    """
    frame = frame - frame.mean()
    if float(np.dot(frame, frame)) < 1e-6:
        return None
    corr = np.correlate(frame, frame, mode="full")[len(frame) - 1:]
    corr /= corr[0] + 1e-12
    lo = int(SAMPLE_RATE / F0_MAX)
    hi = min(int(SAMPLE_RATE / F0_MIN), len(corr) - 1)
    if hi <= lo:
        return None
    peak = int(np.argmax(corr[lo:hi])) + lo
    return SAMPLE_RATE / peak if corr[peak] >= F0_VOICED else None


def segment_pitch(samples: np.ndarray) -> float | None:
    """Частота голоса на отрезке: медиана по озвученным кадрам.

    Громкость и спектр у собеседников в одной комнате похожи - на этом
    разделение по голосам и провалилось. Частота зависит от говорящего,
    а не от его места за столом, поэтому разделяет там, где остальное нет.
    """
    size = int(F0_FRAME_SEC * SAMPLE_RATE)
    hop = int(F0_HOP_SEC * SAMPLE_RATE)
    values = [f for start in range(0, max(0, len(samples) - size), hop)
              if (f := frame_pitch(samples[start:start + size])) is not None]
    return round(float(np.median(values)), 1) if len(values) >= 3 else None


def pitch_split(turns: list[dict]) -> dict:
    """Расходятся ли голоса по частоте настолько, чтобы на это опираться.

    Признак работает только в парах разного пола. В однополых парах разрыв
    в пределах шума, и подсказывать по нему нельзя - собьёт.
    """
    values = sorted(t["acoustics"]["pitch_hz"] for t in turns
                    if (t.get("acoustics") or {}).get("pitch_hz"))
    if len(values) < 8:
        return {"pitch_separates": False}
    middle = len(values) // 2
    low = float(np.median(values[:middle]))
    high = float(np.median(values[middle:]))
    return {
        "pitch_separates": bool(high - low >= PITCH_SPLIT_HZ),
        "pitch_low_hz": round(low, 1),
        "pitch_high_hz": round(high, 1),
    }


def segment_acoustics(samples: np.ndarray) -> dict:
    """Считает признаки, связанные с расстоянием до микрофона."""
    if len(samples) < SAMPLE_RATE // 4:
        return {}

    rms = float(np.sqrt((samples ** 2).mean()) + 1e-9)
    spectrum = np.abs(np.fft.rfft(samples * np.hanning(len(samples))))
    freqs = np.fft.rfftfreq(len(samples), 1 / SAMPLE_RATE)
    total = float(spectrum.sum()) + 1e-9

    return {
        # Громкость: ближний к микрофону обычно громче
        "loudness_db": round(20 * float(np.log10(rms)), 1),
        # Доля высоких частот: у дальнего голоса верх съедается расстоянием
        "hf_ratio": round(float(spectrum[(freqs > 3000) & (freqs < 7000)].sum()) / total, 3),
        # Центр тяжести спектра, тот же смысл другим числом
        "centroid_hz": round(float((spectrum * freqs).sum()) / total),
        # Единственный признак, не зависящий от места за столом.
        "pitch_hz": segment_pitch(samples),
    }


def separation_quality(embeddings: np.ndarray, labels: np.ndarray) -> float:
    """Насколько уверенно голоса разошлись по кластерам.

    Единица - участники различаются чётко, около нуля - модель фактически
    не видит разницы и разметка по говорящим случайна. На записях с одним
    микрофоном в комнате значение обычно низкое.
    """
    if len(embeddings) < 4 or len(set(labels.tolist())) < 2:
        return 0.0

    norm = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-9)
    distances = 1 - norm @ norm.T
    scores = []
    for i in range(len(norm)):
        own = distances[i][(labels == labels[i])]
        other = distances[i][(labels != labels[i])]
        own = own[own > 0]
        if len(own) and len(other):
            a, b = float(own.mean()), float(other.mean())
            scores.append((b - a) / max(a, b))
    return round(float(np.mean(scores)), 3) if scores else 0.0