app/services/timeline.py

"""Whole-track scan.

The opening clip gives AURIS its verdict. This module reads the rest of the
track: the first six minutes are cut into windows of about 30 s, every window
goes through the same features and the same LightGBM model, and the readings
come back as a timeline the page draws under the waveform.

A window score is the classifier applied to an excerpt shorter than the 60 s
clips it was trained on, and every training clip was either all AI or all
human. A timeline therefore shows where the music reads AI-like, not where
an AI took over.
"""

from __future__ import annotations

import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np

from .feature_extractor import extract_features_array
from .inference_xai import XAIInferenceService
from .logging_config import get_logger
from .vocal_analyzer import analyze_vocals_array

logger = get_logger(__name__)

TARGET_WINDOW_SEC = 30.0
# A window quieter than this never reaches the classifier.
SILENCE_PEAK = 1e-3
PEAK_POINTS = 600
# Windows not started within this many seconds are reported as skipped.
SCAN_BUDGET_SEC = float(os.getenv("AURIS_SCAN_BUDGET_SEC", "180"))
SCAN_WORKERS = max(1, int(os.getenv("AURIS_SCAN_WORKERS", "2")))


def plan_windows(duration: float) -> List[Tuple[float, float]]:
    """Equal windows of about 30 s that tile the whole duration (20 to 45 s each)."""
    if duration <= 0:
        return []
    count = max(1, int(round(duration / TARGET_WINDOW_SEC)))
    size = duration / count
    return [
        (round(i * size, 3), round(duration if i == count - 1 else (i + 1) * size, 3))
        for i in range(count)
    ]


def waveform_peaks(y: np.ndarray, points: int = PEAK_POINTS) -> List[float]:
    """Peak level per slice of the track, 0 to 1, for drawing the waveform."""
    if y.size == 0:
        return []
    edges = np.linspace(0, y.size, min(points, y.size) + 1).astype(int)
    peaks = np.array([
        float(np.max(np.abs(y[a:b]))) if b > a else 0.0
        for a, b in zip(edges[:-1], edges[1:])
    ])
    top = float(peaks.max())
    if top <= 0:
        return [0.0] * len(peaks)
    return [round(float(v), 3) for v in peaks / top]


def _read_window(
    y: np.ndarray,
    sr: int,
    index: int,
    start: float,
    end: float,
    xai: XAIInferenceService,
) -> Dict[str, Any]:
    base: Dict[str, Any] = {"index": index, "start": round(start, 2), "end": round(end, 2)}
    seg = np.ascontiguousarray(y[int(start * sr):min(y.size, int(end * sr))])
    if seg.size < sr:
        return {**base, "state": "skipped"}
    if float(np.max(np.abs(seg))) < SILENCE_PEAK:
        return {**base, "state": "silent"}
    try:
        features = extract_features_array(seg, sr)
        vocals = None
        try:
            vocals = analyze_vocals_array(seg, sr)
        except Exception as e:  # noqa: BLE001 - the main pipeline also goes on without vocals
            logger.warning(f"Window {index}: vocal analysis failed: {e}")
        reading = xai.predict_window(features, vocals, window_sec=end - start)
        if reading is None:
            return {**base, "state": "failed"}
        return {
            **base,
            "state": "ok",
            "probability": round(reading.probability, 4),
            "isAi": reading.is_ai,
            "margin": reading.margin,
            "hasVocals": bool(vocals and vocals.has_vocals),
            "reasons": [
                {
                    "name": c.name,
                    "label": c.label,
                    "labelEn": c.label_en,
                    "direction": c.direction,
                    "shapValue": round(c.shap_value, 4),
                }
                for c in reading.reasons
            ],
        }
    except Exception as e:  # noqa: BLE001 - one bad window must not sink the scan
        logger.warning(f"Window {index} ({start:.0f}-{end:.0f}s) failed: {e}")
        return {**base, "state": "failed"}


def _summarise(segments: List[Dict[str, Any]]) -> Dict[str, Any]:
    scored = [s for s in segments if s["state"] == "ok"]
    seconds = sum(s["end"] - s["start"] for s in scored)
    if not scored or seconds <= 0:
        return {
            "scoredCount": 0, "flaggedCount": 0, "scoredSec": 0.0,
            "aiShare": 0.0, "meanProbability": 0.0, "maxProbability": 0.0, "peakIndex": None,
        }
    flagged = [s for s in scored if s["isAi"]]
    peak = max(scored, key=lambda s: s["probability"])
    return {
        "scoredCount": len(scored),
        "flaggedCount": len(flagged),
        "scoredSec": round(seconds, 1),
        "aiShare": round(sum(s["end"] - s["start"] for s in flagged) / seconds, 3),
        "meanProbability": round(sum(s["probability"] * (s["end"] - s["start"]) for s in scored) / seconds, 4),
        "maxProbability": peak["probability"],
        "peakIndex": peak["index"],
    }


def scan_track(
    y: np.ndarray,
    sr: int,
    xai: XAIInferenceService,
    *,
    total_sec: float,
    on_progress: Optional[Callable[[int, int], None]] = None,
    cancelled: Optional[Callable[[], bool]] = None,
    budget_sec: float = SCAN_BUDGET_SEC,
) -> Dict[str, Any]:
    """Read every window of `y` and return the timeline payload (blocking)."""
    duration = y.size / sr
    windows = plan_windows(duration)
    started = time.monotonic()
    lock = threading.Lock()
    finished = 0

    def work(item: Tuple[int, Tuple[float, float]]) -> Dict[str, Any]:
        nonlocal finished
        index, (start, end) = item
        if (cancelled is not None and cancelled()) or time.monotonic() - started > budget_sec:
            out: Dict[str, Any] = {"index": index, "start": round(start, 2), "end": round(end, 2), "state": "skipped"}
        else:
            out = _read_window(y, sr, index, start, end, xai)
        with lock:
            finished += 1
            if on_progress is not None:
                on_progress(finished, len(windows))
        return out

    with ThreadPoolExecutor(max_workers=min(SCAN_WORKERS, max(1, len(windows)))) as pool:
        segments = list(pool.map(work, enumerate(windows)))

    logger.info(
        f"Timeline: {len(windows)} windows over {duration:.0f}s in {time.monotonic() - started:.1f}s, "
        f"{sum(1 for s in segments if s['state'] == 'ok')} scored"
    )
    return {
        "durationSec": round(duration, 2),
        "totalSec": round(max(total_sec, duration), 2),
        "truncated": total_sec - duration > 1.0,
        "threshold": round(float(xai.threshold), 6),
        "peaks": waveform_peaks(y),
        "segments": segments,
        "summary": _summarise(segments),
    }