"""Подбор оптимального числа потоков на конкретной машине.

Масштабирование ONNX сильно зависит от процессора: на Apple M4 оптимум
оказался в районе 4 потоков, на Ryzen картина другая. Гадать бесполезно,
поэтому сервис умеет померить сам.
"""
import time
from pathlib import Path

import numpy as np

from app.pipeline import SAMPLE_RATE, Pipeline

__all__ = ["run_benchmark"]

# Больше 16 потоков смысла мерить нет: везде, где замеряли, деградация
# начиналась гораздо раньше.
THREAD_OPTIONS = (1, 2, 4, 8, 16)


def run_benchmark(samples: np.ndarray, models_dir: Path, base_dir: Path,
                  replacements_path: Path, seconds: float = 60.0,
                  speakers: int = 2) -> dict:
    """Гоняет один и тот же фрагмент на разном числе потоков."""
    clip = samples[: int(seconds * SAMPLE_RATE)]
    duration = len(clip) / SAMPLE_RATE
    rows = []

    for threads in THREAD_OPTIONS:
        pipe = Pipeline(models_dir=models_dir, threads=threads,
                        replacements_path=replacements_path, base_dir=base_dir)
        pipe.warmup()

        t0 = time.time()
        pipe._diarizer(speakers).process(clip)
        diar = time.time() - t0

        t0 = time.time()
        pipe._asr.recognize(clip[: int(30 * SAMPLE_RATE)], sample_rate=SAMPLE_RATE)
        asr = time.time() - t0

        rows.append({
            "threads": threads,
            "diarization_sec": round(diar, 2),
            "diarization_x": round(duration / diar, 1) if diar else 0,
            "asr_sec": round(asr, 2),
            "asr_x": round(min(30.0, duration) / asr, 1) if asr else 0,
        })

    best_diar = max(rows, key=lambda r: r["diarization_x"])
    best_asr = max(rows, key=lambda r: r["asr_x"])
    return {
        "clip_sec": round(duration, 1),
        "results": rows,
        "best_for_diarization": best_diar["threads"],
        "best_for_asr": best_asr["threads"],
        "recommended_threads": best_diar["threads"],
        "hint": (
            f"Поставьте threads = {best_diar['threads']} в config.toml. "
            "Разделение говорящих занимает большую часть времени, поэтому "
            "ориентируемся на него."
        ),
    }