#!/usr/bin/env python3
"""Rank music candidates against the voice they will sit under.

A track is not chosen because its tags matched the query. The thing that decides
whether music works under speech is measurable, and none of it is in the
metadata:

  * **mid occupancy** — how much of the track's energy sits in 300–3400 Hz, the
    band the voice needs. A pad-and-sub track leaves that band empty and can be
    played 6 dB louder than a guitar track at the same LUFS while still being
    fully intelligible under speech. This is the single strongest predictor and
    it is why "the music is barely audible" and "the music fights the voice" are
    the same problem: a mid-heavy track has to be ducked into inaudibility.
  * **syllabic modulation** — energy in the 1–4 kHz band modulated at 2–8 Hz.
    That is the rate human speech opens and closes at, so a track with a strong
    figure there reads as a second person talking. It is an indicator, not a
    verdict: report the number, let the editor listen.
  * **onset density and tempo** — whether the track has a pulse the edit can cut
    against at all.
  * **intro ramp** — how long before the track reaches full energy. A track that
    starts at full tilt cannot open a video under a first sentence.
  * **dynamic spread** — a heavily limited track flattens the whole mix; a very
    dynamic one keeps ducking in and out under the sidechain.

Every candidate also gets a 12-second excerpt written next to the report, taken
from its own loudest passage, because a track cannot be chosen from numbers
alone either.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))

from skill_config import emit, require_binary, utf8_stdout  # noqa: E402

RATE = 22050
FRAME = 2048
HOP = 512


# Сколько секунд трека достаточно, чтобы понять его характер.
#
# Не косметическое ограничение. Спектрограмма целиком строится в памяти, и на
# шестиминутном треке это 255 тысяч кадров на 1025 бинов в complex128 — 3.9 ГБ,
# после чего numpy падает. Полторы минуты описывают трек не хуже: доля энергии в
# речевой полосе, плотность атак сверху и темп на втором припеве те же, что на
# первом. Интро при этом попадает целиком, а оно и есть единственное, что важно
# знать про начало.
ANALYSIS_SECONDS = 100


def decode(path: Path) -> np.ndarray:
    done = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-i", str(path),
         "-t", str(ANALYSIS_SECONDS),
         "-map", "0:a:0", "-f", "f32le", "-ac", "1", "-ar", str(RATE), "-"],
        capture_output=True)
    if done.returncode:
        raise SystemExit((done.stderr or b"").decode("utf-8", "replace")[-800:])
    return np.frombuffer(done.stdout, dtype="<f4").astype(np.float64)


def spectrogram(signal: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    frames = 1 + max(0, (signal.size - FRAME)) // HOP
    if frames < 8:
        raise SystemExit("Track too short to analyse")
    view = np.lib.stride_tricks.as_strided(
        signal, shape=(frames, FRAME),
        strides=(signal.strides[0] * HOP, signal.strides[0]))
    magnitude = np.abs(np.fft.rfft(view * np.hanning(FRAME), axis=1))
    return magnitude, np.fft.rfftfreq(FRAME, 1.0 / RATE)


def true_duration(path: Path) -> float:
    done = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=nw=1:nk=1", str(path)],
        capture_output=True, text=True, errors="replace")
    try:
        return float((done.stdout or "0").strip())
    except ValueError:
        return 0.0


def analyse(path: Path) -> dict:
    signal = decode(path)
    # Анализируется первая сотня секунд, но длина трека нужна настоящая: по ней
    # решается, придётся ли его зацикливать под ролик.
    analysed = signal.size / RATE
    duration = true_duration(path) or analysed
    magnitude, freqs = spectrogram(signal)
    power = magnitude ** 2
    frame_rate = RATE / HOP

    total = power.sum(axis=1) + 1e-20
    mid = power[:, (freqs >= 300) & (freqs < 3400)].sum(axis=1)
    low = power[:, freqs < 250].sum(axis=1)
    high = power[:, freqs >= 5000].sum(axis=1)

    mid_occupancy = float(np.median(mid / total))
    low_share = float(np.median(low / total))
    high_share = float(np.median(high / total))

    # How much of the track's 1–4 kHz envelope moves at syllable rate (2–8 Hz).
    # Comparing band *sums* here would say nothing: 2–8 Hz is 6 Hz of a 19.8 Hz
    # reference, so a completely flat modulation spectrum already scores 0.30 and
    # every track lands near the threshold. Comparing densities makes 1.0 mean
    # "no more active at speech rate than anywhere else" and the number readable.
    voice_band = np.sqrt(power[:, (freqs >= 1000) & (freqs < 4000)].sum(axis=1))
    envelope = voice_band - voice_band.mean()
    spectrum = np.abs(np.fft.rfft(envelope * np.hanning(envelope.size)))
    mod_freqs = np.fft.rfftfreq(envelope.size, 1.0 / frame_rate)
    inner = (mod_freqs >= 2.0) & (mod_freqs <= 8.0)
    outer = (mod_freqs >= 0.2) & (mod_freqs <= 20.0)
    syllabic_ratio = float((spectrum[inner].mean() if inner.any() else 0.0)
                           / ((spectrum[outer].mean() if outer.any() else 1e-9) + 1e-9))

    # Sharp transients above 5 kHz — hi-hats, rimshots, ride bells. This is the
    # one metric that explains the complaint "the music is inaudible and its top
    # end still pokes through the voice". A three-band ducker holds the high band
    # back least, because air and cymbals sit above speech and normally do not
    # mask it. That reasoning holds for sustained top end and fails completely
    # for percussive top end: a hi-hat is a click, it survives any duck with a
    # release long enough not to pump, and it is the only part of the track the
    # viewer hears. Meanwhile the body of the music — which is what "present"
    # actually sounds like — is fully ducked away.
    high_band = np.sqrt(power[:, freqs >= 5000].sum(axis=1))
    high_flux = np.maximum(0.0, np.diff(high_band))
    hf_transient = float(
        # Делить надо на длину **проанализированного** куска, а не всего трека:
        # поток посчитан только по нему, и деление на полную длину занизило бы
        # плотность вдвое-втрое на длинных записях.
        (high_flux > high_flux.mean() + 2.0 * high_flux.std()).sum() / max(1.0, analysed))

    # Onsets and tempo from the spectral flux.
    flux = np.maximum(0.0, np.diff(np.sqrt(power).sum(axis=1)))
    flux = flux - flux.mean()
    if flux.size > 32:
        correlation = np.correlate(flux, flux, mode="full")[flux.size - 1:]
        lo, hi = int(frame_rate * 60 / 180), int(frame_rate * 60 / 60)
        window = correlation[lo:hi]
        bpm = float(60.0 * frame_rate / (lo + int(np.argmax(window)))) if window.size else 0.0
        pulse = float(window.max() / (correlation[0] + 1e-9)) if window.size else 0.0
    else:
        bpm, pulse = 0.0, 0.0

    level_db = 20.0 * np.log10(np.sqrt(total / FRAME) + 1e-12)
    smooth = np.convolve(level_db, np.ones(int(frame_rate)) / frame_rate, mode="same")
    reference_level = float(np.percentile(smooth, 60))
    ramp_idx = np.where(smooth >= reference_level - 3.0)[0]
    intro_seconds = round(float(ramp_idx[0] / frame_rate), 2) if ramp_idx.size else 0.0
    spread = float(np.percentile(smooth, 95) - np.percentile(smooth, 20))

    return {
        "file": str(path),
        "duration": round(duration, 2),
        "mid_occupancy": round(mid_occupancy, 3),
        "low_share": round(low_share, 3),
        "high_share": round(high_share, 3),
        "syllabic_ratio": round(syllabic_ratio, 3),
        "hf_transients_per_second": round(hf_transient, 2),
        "bpm": round(bpm, 1),
        "pulse_strength": round(pulse, 3),
        "intro_seconds": intro_seconds,
        "dynamic_spread_db": round(spread, 1),
        "peak_dbfs": round(float(20 * np.log10(np.max(np.abs(signal)) + 1e-12)), 2),
        "_loud_at": float(int(np.argmax(smooth)) / frame_rate),
    }


def score(measured: dict, brief: dict) -> dict:
    """Weighted fit. Every component says what it punished and by how much."""
    notes: list[str] = []
    points = 100.0

    occupancy = measured["mid_occupancy"]
    if occupancy > 0.34:
        penalty = min(40.0, (occupancy - 0.34) * 260)
        points -= penalty
        notes.append(f"{occupancy:.0%} энергии в речевой полосе — трек придётся задавить, "
                     f"чтобы речь читалась (−{penalty:.0f})")
    else:
        notes.append(f"речевая полоса свободна ({occupancy:.0%}) — трек будет слышен под голосом")

    syllabic = measured["syllabic_ratio"]
    if syllabic > 1.6:
        penalty = min(22.0, (syllabic - 1.6) * 22)
        points -= penalty
        notes.append(f"активность речевой полосы на слоговом ритме ×{syllabic:.2f} — "
                     f"вокал или солирующая линия маскирует согласные (−{penalty:.0f})")
    else:
        notes.append(f"на слоговом ритме спокойно (×{syllabic:.2f})")

    hf = measured["hf_transients_per_second"]
    if hf > 2.5:
        penalty = min(26.0, (hf - 2.5) * 7.0)
        points -= penalty
        notes.append(f"{hf:.1f} резких транзиента в секунду выше 5 кГц — хэты пробьют "
                     f"любой дакинг и останутся единственным, что слышно из трека "
                     f"(−{penalty:.0f})")
    else:
        notes.append(f"верх спокойный ({hf:.1f} транзиента/с) — трек будет слышен телом, "
                     "а не щелчками")

    if measured["peak_dbfs"] > 0.0:
        penalty = min(14.0, 4 + measured["peak_dbfs"] * 3)
        points -= penalty
        notes.append(f"пик {measured['peak_dbfs']:+.1f} дБFS — исходник уже клиппован, "
                     f"чище он не станет (−{penalty:.0f})")

    low, high = brief.get("bpm_range", (0, 0))
    if low and measured["bpm"] and not (low <= measured["bpm"] <= high):
        distance = min(abs(measured["bpm"] - low), abs(measured["bpm"] - high))
        penalty = min(15.0, distance * 0.25)
        points -= penalty
        notes.append(f"{measured['bpm']:.0f} BPM вне запрошенных {low}–{high} (−{penalty:.0f})")

    if measured["pulse_strength"] < 0.12:
        points -= 10
        notes.append(f"пульс слабый ({measured['pulse_strength']:.2f}) — резать по битам не выйдет")

    if measured["intro_seconds"] > 12:
        points -= 8
        notes.append(f"интро {measured['intro_seconds']} с — слишком долго для короткого ролика")
    elif measured["intro_seconds"] < 0.5:
        points -= 5
        notes.append("трек стартует сразу на полной энергии — под первую фразу не подложить")

    if measured["dynamic_spread_db"] < 4:
        points -= 6
        notes.append(f"разброс {measured['dynamic_spread_db']} дБ — трек закатан в кирпич, "
                     "будет давить на голос ровной стеной")
    elif measured["dynamic_spread_db"] > 20:
        points -= 6
        notes.append(f"разброс {measured['dynamic_spread_db']} дБ — под сайдчейном будет нырять")

    if measured["duration"] < brief.get("min_duration", 0):
        points -= 12
        notes.append(f"{measured['duration']} с короче ролика — придётся зацикливать")

    return {"score": round(max(0.0, points), 1), "notes": notes}


BRIEFS: dict[str, dict] = {
    "calm_talk": {"bpm_range": (70, 105), "use": "Разговорное видео, эксперт, объяснение."},
    "energetic": {"bpm_range": (110, 150), "use": "Клиповый монтаж, продукт, реклама."},
    "cinematic": {"bpm_range": (60, 95), "use": "Атмосфера, сторителлинг, финал."},
    "drive": {"bpm_range": (120, 175), "use": "Спорт, мото, экшн."},
}


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("candidates", nargs="+", help="Audio files to rank")
    parser.add_argument("--brief", default="calm_talk", choices=list(BRIEFS))
    parser.add_argument("--under-duration", type=float, default=0.0,
                        help="Length of the video the track has to cover")
    parser.add_argument("--excerpt-dir", help="Write a 12 s excerpt per candidate here")
    parser.add_argument("--out", help="Write the ranked report as JSON")
    args = parser.parse_args()

    require_binary("ffmpeg")
    brief = dict(BRIEFS[args.brief])
    brief["min_duration"] = args.under_duration

    ranked = []
    for name in args.candidates:
        path = Path(name).resolve()
        if not path.is_file():
            ranked.append({"file": str(path), "error": "not found"})
            continue
        measured = analyse(path)
        loud_at = measured.pop("_loud_at")
        verdict = score(measured, brief)
        entry = {**measured, **verdict}
        if args.excerpt_dir:
            out_dir = Path(args.excerpt_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
            excerpt = out_dir / f"{path.stem}_excerpt.mp3"
            start = max(0.0, min(loud_at - 4.0, max(0.0, measured["duration"] - 12.0)))
            subprocess.run(
                ["ffmpeg", "-y", "-v", "error", "-nostdin", "-ss", f"{start:.2f}",
                 "-i", str(path), "-t", "12", "-af",
                 "afade=t=in:d=0.4,afade=t=out:st=11.4:d=0.6",
                 "-c:a", "libmp3lame", "-b:a", "160k", str(excerpt)],
                capture_output=True)
            entry["excerpt"] = str(excerpt)
            entry["excerpt_from_seconds"] = round(start, 2)
        ranked.append(entry)

    ranked.sort(key=lambda item: item.get("score", -1), reverse=True)
    payload = {
        "status": "ok",
        "brief": args.brief,
        "brief_use": BRIEFS[args.brief]["use"],
        "ranked": ranked,
        "how_to_read": {
            "mid_occupancy": "Доля энергии в 300–3400 Гц. Ниже 0.30 — трек не спорит с голосом; "
                             "выше 0.40 — под речью его будет не слышно, сколько ни поднимай.",
            "syllabic_ratio": "Плотность модуляции 2–8 Гц против всего диапазона 0.2–20 Гц в "
                              "полосе 1–4 кГц. 1.0 — ровно; выше 1.6 — в речевой полосе что-то "
                              "движется со скоростью речи (вокал, солирующая линия).",
            "peak_dbfs": "Выше 0 — исходник клиппован ещё до монтажа.",
            "hf_transients_per_second": "Резкие атаки выше 5 кГц. Выше 2.5 — хэты и "
                                        "римшоты пробьют дакинг и станут единственным, "
                                        "что слышно из трека, пока его тело задавлено.",
            "pulse_strength": "Насколько выражен бит. Ниже 0.12 — по нему не порезать.",
        },
        "next": ("Послушай excerpt у двух-трёх верхних кандидатов и только потом бери трек. "
                 "Число ранжирует, но не выбирает."),
    }
    if args.out:
        Path(args.out).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        payload["report"] = args.out
    emit(payload)


if __name__ == "__main__":
    main()
