analyzer.py

"""
OpenMix audio track analysis: tempo, key, vocals, energy, spectral features.
Pure analysis logic — no mixing/crossfade state.
"""

import logging
from pathlib import Path
from typing import List, Optional, Tuple

import librosa
import numpy as np
from scipy import signal

from audio_utils import analysis_mono, normalize_audio
from models import AudioConfig, TrackAnalysis

logger = logging.getLogger(__name__)


def detect_silence(y: np.ndarray, sr: int, threshold_db: float = -50.0, min_duration: float = 0.5) -> Tuple[float, float]:
    """Detect leading and trailing silence. Returns (trim_start, trim_end) in seconds."""
    if y.ndim > 1:
        y_mono = np.mean(y, axis=1)
    else:
        y_mono = y

    frame_length = 2048
    hop_length = 512
    rms = librosa.feature.rms(y=y_mono, frame_length=frame_length, hop_length=hop_length)[0]

    threshold = 10 ** (threshold_db / 20)
    silent_frames = rms < threshold

    if not np.any(silent_frames):
        return 0.0, 0.0

    # Find first non-silent frame
    first_sound = np.argmax(~silent_frames)
    trim_start = first_sound * hop_length / sr

    # Find last non-silent frame
    last_sound = len(silent_frames) - 1 - np.argmax(silent_frames[::-1])
    trim_end = last_sound * hop_length / sr

    # Only trim if silence exceeds minimum duration
    if trim_start < min_duration:
        trim_start = 0.0
    if len(y) / sr - trim_end < min_duration:
        trim_end = len(y) / sr

    return trim_start, trim_end


def analyze_track(file_path: Path, config: AudioConfig) -> Optional[TrackAnalysis]:
    """Analyze a single audio file and return full feature set."""
    try:
        logger.info(f"Analyzing: {file_path.name}")

        y, sr = librosa.load(str(file_path), sr=config.sample_rate, mono=False)
        if y.ndim > 1:
            y = y.T

        y_normalized = normalize_audio(y)
        mono = analysis_mono(y_normalized)

        # Detect and trim leading/trailing silence
        trim_start, trim_end = detect_silence(y, sr)
        if trim_start > 0 or trim_end < len(y) / sr:
            start_sample = int(trim_start * sr)
            end_sample = int(trim_end * sr)
            if end_sample > start_sample:
                y = y[start_sample:end_sample]
                mono = analysis_mono(normalize_audio(y))
                logger.info(f"    Trimmed silence: {trim_start:.1f}s - {trim_end:.1f}s")

        max_samples = int(sr * config.max_analysis_seconds)
        focus = mono[:max_samples] if len(mono) > max_samples else mono

        if len(focus) < 2048:
            logger.warning(f"{file_path.name}: very short signal, using fallback")
            tempo = 120.0
            beats = np.array([])
            beat_frames = np.array([])
        else:
            tempo, beat_frames = librosa.beat.beat_track(y=focus, sr=sr, units='frames')
            tempo = float(np.asarray(tempo).reshape(-1)[0])
            beats = librosa.frames_to_time(beat_frames, sr=sr, hop_length=512)

        n_fft = max(2, min(2048, len(focus)))
        hop_length = max(1, min(512, max(1, n_fft // 4)))

        chroma = librosa.feature.chroma_stft(y=focus, sr=sr, n_fft=n_fft, hop_length=hop_length)
        key_profile = np.mean(chroma, axis=1)
        key = int(np.argmax(key_profile))

        rms = librosa.feature.rms(y=focus, frame_length=n_fft, hop_length=hop_length)[0]
        energy = float(np.mean(rms))
        energy_variation = float(np.std(rms))

        spectral_centroid = float(np.mean(
            librosa.feature.spectral_centroid(y=focus, sr=sr, n_fft=n_fft, hop_length=hop_length)
        ))
        spectral_rolloff = float(np.mean(
            librosa.feature.spectral_rolloff(y=focus, sr=sr, n_fft=n_fft, hop_length=hop_length)
        ))
        spectral_bandwidth = float(np.mean(
            librosa.feature.spectral_bandwidth(y=focus, sr=sr, n_fft=n_fft, hop_length=hop_length)
        ))
        zcr = float(np.mean(librosa.feature.zero_crossing_rate(focus)))

        vocal_segments = detect_vocals_edges(mono, sr, config)
        intro_end, outro_start = detect_intro_outro(focus, sr, beats)

        peak_level = float(np.max(np.abs(y_normalized)))
        rms_level = float(np.sqrt(np.mean(y_normalized**2)))

        return TrackAnalysis(
            file_path=file_path,
            duration=len(y) / sr,
            tempo=tempo,
            beats=beats,
            beat_frames=beat_frames,
            key=key,
            energy=energy,
            energy_variation=energy_variation,
            spectral_centroid=spectral_centroid,
            spectral_rolloff=spectral_rolloff,
            spectral_bandwidth=spectral_bandwidth,
            zcr=zcr,
            vocal_segments=vocal_segments,
            intro_end=intro_end,
            outro_start=outro_start,
            peak_level=peak_level,
            rms_level=rms_level,
            audio_data=y,
            sample_rate=sr,
        )

    except Exception as e:
        logger.error(f"Error analyzing {file_path.name}: {e}")
        return None


def detect_vocals(y: np.ndarray, sr: int) -> List[Tuple[float, float]]:
    """Detect vocal segments using harmonic-percussive separation."""
    try:
        if len(y) < 2048:
            return []

        n_fft = max(2, min(2048, len(y)))
        hop_length = max(1, min(512, max(1, n_fft // 4)))

        y_harmonic = librosa.effects.hpss(y)[0]

        spec_centroid = librosa.feature.spectral_centroid(y=y_harmonic, sr=sr, n_fft=n_fft, hop_length=hop_length)[0]
        chroma = librosa.feature.chroma_stft(y=y_harmonic, sr=sr, n_fft=n_fft, hop_length=hop_length)
        chroma_strength = np.sum(chroma, axis=0)

        # Harmonic content ratio: vocals have higher harmonic energy than percussion
        harmonic_energy = np.mean(y_harmonic ** 2)
        total_energy = np.mean(y ** 2) + 1e-10
        harmonic_ratio = harmonic_energy / total_energy

        # If harmonic content is low, likely percussion not vocals
        if harmonic_ratio < 0.3:
            logger.info(f"    Low harmonic ratio ({harmonic_ratio:.2f}), skipping vocal detection")
            return []

        spec_norm = (spec_centroid - np.mean(spec_centroid)) / (np.std(spec_centroid) + 1e-8)
        chroma_norm = (chroma_strength - np.mean(chroma_strength)) / (np.std(chroma_strength) + 1e-8)

        vocal_prob = (np.clip(spec_norm * 0.4, -1, 1) + np.clip(chroma_norm * 0.4, -1, 1)) / 2.0

        if len(vocal_prob) > 10:
            window_size = min(21, len(vocal_prob) // 5)
            if window_size >= 5:
                vocal_prob = signal.savgol_filter(vocal_prob, window_size | 1, 2)

        frame_times = librosa.frames_to_time(np.arange(len(vocal_prob)), sr=sr, hop_length=hop_length)

        vocal_threshold = 0.2
        vocal_frames = vocal_prob > vocal_threshold

        segments = []
        in_vocal = False
        start_time = 0.0

        for i, is_vocal in enumerate(vocal_frames):
            current_time = frame_times[i] if i < len(frame_times) else frame_times[-1]
            if is_vocal and not in_vocal:
                start_time = current_time
                in_vocal = True
            elif not is_vocal and in_vocal:
                if current_time - start_time > 2.0:
                    segments.append((start_time, current_time))
                in_vocal = False

        if in_vocal and len(frame_times) > 0:
            if frame_times[-1] - start_time > 2.0:
                segments.append((start_time, frame_times[-1]))

        logger.info(f"    Detected {len(segments)} vocal segments")
        return segments

    except Exception as e:
        logger.warning(f"    Vocal detection failed: {e}, treating as no vocals")
        return []


def detect_vocals_edges(y: np.ndarray, sr: int, config: AudioConfig) -> List[Tuple[float, float]]:
    """Detect vocals on intro/outro windows only (faster for long tracks)."""
    if len(y) < 2048:
        return []

    scan_samples = int(sr * config.vocal_scan_seconds)

    if len(y) <= scan_samples * 2:
        return detect_vocals(y, sr)

    intro_audio = y[:scan_samples]
    outro_audio = y[-scan_samples:]
    outro_offset = (len(y) - scan_samples) / sr

    intro_segments = detect_vocals(intro_audio, sr)
    outro_local = detect_vocals(outro_audio, sr)
    outro_segments = [(s + outro_offset, e + outro_offset) for s, e in outro_local]

    merged = intro_segments + outro_segments
    logger.info(f"    Detected {len(merged)} vocal segments (intro/outro scan)")
    return merged


def detect_intro_outro(y: np.ndarray, sr: int, beats: np.ndarray) -> Tuple[float, float]:
    """Detect intro and outro sections using RMS energy profile."""
    duration = len(y) / sr

    n_fft = max(2, min(2048, len(y)))
    hop_length = max(1, min(512, max(1, n_fft // 4)))
    rms = librosa.feature.rms(y=y, frame_length=n_fft, hop_length=hop_length)[0]
    hop_time = hop_length / sr
    time_frames = np.arange(len(rms)) * hop_time

    energy_threshold = (np.max(rms) - np.min(rms)) * 0.15 + np.min(rms)
    above_threshold = rms > energy_threshold

    if np.any(above_threshold):
        intro_end = time_frames[np.argmax(above_threshold)]
        outro_start = time_frames[len(rms) - 1 - np.argmax(above_threshold[::-1])]
    else:
        intro_end = duration * 0.15
        outro_start = duration * 0.85

    # Minimum intro/outro duration: 5% of track
    min_section = duration * 0.05
    intro_end = max(intro_end, min_section)
    outro_start = min(outro_start, duration - min_section)

    if len(beats) > 32:
        beat_intro_end = beats[min(16, len(beats) // 4)]
        beat_outro_start = beats[max(-16, -len(beats) // 4)]
        intro_end = min(intro_end, beat_intro_end)
        outro_start = max(outro_start, beat_outro_start)

    return intro_end, outro_start