audio_analysis.py

from __future__ import annotations

import subprocess
from pathlib import Path

import numpy as np

PITCH_CLASSES = ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B")

# Recognized tala / meter cycles keyed by beats-per-cycle.
TALA_SIGNATURES: dict[int, tuple[str, str]] = {
    3: ("Tisra Eka", "3/4"),
    4: ("Common Time", "4/4"),
    5: ("Khanda Chapu", "5/4"),
    6: ("Rupaka", "6/8"),
    7: ("Misra Chapu", "7/8"),
    8: ("Adi Tala", "4/4"),
    9: ("Sankirna Chapu", "9/8"),
    10: ("Jhaptaal", "5/4"),
    12: ("Ektaal", "12/8"),
    14: ("Dhamar", "7/4"),
    16: ("Teentaal", "4/4"),
}

# Krumhansl-Schmuckler key profiles.
_MAJOR_PROFILE = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88], dtype=np.float64)
_MINOR_PROFILE = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17], dtype=np.float64)


def _safe_normalize(x: np.ndarray) -> np.ndarray:
    peak = float(np.max(np.abs(x))) if x.size else 0.0
    if peak <= 1e-9:
        return x.astype(np.float32)
    return (x / peak).astype(np.float32)


def _decode_audio_mono(audio_path: Path, sample_rate: int = 22050) -> np.ndarray:
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(audio_path),
        "-f",
        "f32le",
        "-ac",
        "1",
        "-ar",
        str(sample_rate),
        "-",
    ]
    proc = subprocess.run(cmd, capture_output=True, check=False)
    if proc.returncode != 0:
        stderr = proc.stderr.decode(errors="replace")[:500]
        raise RuntimeError(f"Audio decode failed: {stderr}")

    y = np.frombuffer(proc.stdout, dtype=np.float32)
    if y.size == 0:
        raise RuntimeError("Audio decode returned empty output")

    y = np.nan_to_num(y)
    return _safe_normalize(y)


def _make_frames(y: np.ndarray, frame_size: int, hop_size: int) -> np.ndarray:
    if y.size < frame_size:
        y = np.pad(y, (0, frame_size - y.size))

    frame_count = 1 + (y.size - frame_size) // hop_size
    if frame_count <= 0:
        return np.empty((0, frame_size), dtype=np.float32)

    idx = np.arange(frame_size)[None, :] + hop_size * np.arange(frame_count)[:, None]
    return y[idx]


def _onset_envelope(y: np.ndarray, frame_size: int = 2048, hop_size: int = 512) -> np.ndarray:
    frames = _make_frames(y, frame_size=frame_size, hop_size=hop_size)
    if frames.shape[0] < 3:
        return np.array([], dtype=np.float32)

    window = np.hanning(frame_size).astype(np.float32)
    spectrum = np.abs(np.fft.rfft(frames * window, axis=1))
    flux = np.maximum(0.0, np.diff(spectrum, axis=0)).sum(axis=1)

    if flux.size == 0:
        return flux.astype(np.float32)

    flux = flux - float(np.min(flux))
    flux = flux / (float(np.max(flux)) + 1e-9)

    # Mild smoothing to reduce spurious peaks.
    kernel = np.ones(5, dtype=np.float32) / 5.0
    flux = np.convolve(flux, kernel, mode="same")
    return flux.astype(np.float32)


def _tempo_from_autocorrelation(onset: np.ndarray, sample_rate: int, hop_size: int) -> tuple[float | None, float]:
    if onset.size < 16:
        return None, 0.0

    x = onset - float(np.mean(onset))
    if np.allclose(x, 0.0, atol=1e-9):
        return None, 0.0

    ac = np.correlate(x, x, mode="full")[x.size - 1 :]
    if ac.size < 4:
        return None, 0.0

    bpm_min, bpm_max = 60.0, 200.0
    lag_min = max(1, int(round((60.0 / bpm_max) * sample_rate / hop_size)))
    lag_max = min(ac.size - 1, int(round((60.0 / bpm_min) * sample_rate / hop_size)))

    if lag_max <= lag_min:
        return None, 0.0

    search = ac[lag_min : lag_max + 1]
    if search.size < 3:
        return None, 0.0

    peak_rel = int(np.argmax(search))
    lag = lag_min + peak_rel

    # Check octave-related tempo aliases (half/double-time).
    candidates = [lag]
    if lag // 2 >= lag_min:
        candidates.append(lag // 2)
    if lag * 2 <= lag_max:
        candidates.append(lag * 2)
    if lag * 3 // 2 <= lag_max:
        candidates.append(lag * 3 // 2)

    candidate_scores: dict[int, float] = {}
    for cand in candidates:
        score = float(ac[cand])
        if cand != lag:
            score *= 0.96
        candidate_scores[cand] = score

    best_lag = max(candidate_scores, key=candidate_scores.get)
    best_score = candidate_scores[best_lag]
    bpm = 60.0 * sample_rate / (best_lag * hop_size)

    # Resolve common half-time/double-time ambiguity.
    if bpm < 85.0 and best_lag // 2 >= lag_min:
        alt_lag = best_lag // 2
        alt_bpm = 60.0 * sample_rate / (alt_lag * hop_size)
        if alt_bpm <= bpm_max and float(ac[alt_lag]) >= best_score * 0.88:
            best_lag = alt_lag
            bpm = alt_bpm
    elif bpm > 170.0 and best_lag * 2 <= lag_max:
        alt_lag = best_lag * 2
        alt_bpm = 60.0 * sample_rate / (alt_lag * hop_size)
        if alt_bpm >= bpm_min and float(ac[alt_lag]) >= best_score * 0.88:
            best_lag = alt_lag
            bpm = alt_bpm

    peak = float(ac[best_lag])
    sorted_peaks = np.sort(search)
    second = float(sorted_peaks[-2]) if sorted_peaks.size > 1 else 0.0
    peak_sep = max(0.0, peak - second) / (abs(peak) + 1e-9)
    signal_strength = max(0.0, peak / (float(np.mean(search)) + 1e-9) - 1.0)
    confidence = float(np.clip(0.65 * peak_sep + 0.35 * (signal_strength / 4.0), 0.0, 1.0))

    if bpm < bpm_min or bpm > bpm_max:
        return None, 0.0

    return float(bpm), confidence


def _tempo_from_frequency(onset: np.ndarray, sample_rate: int, hop_size: int) -> tuple[float | None, float]:
    if onset.size < 32:
        return None, 0.0

    x = onset - float(np.mean(onset))
    if np.allclose(x, 0.0, atol=1e-9):
        return None, 0.0

    spec = np.abs(np.fft.rfft(x * np.hanning(x.size)))
    freqs = np.fft.rfftfreq(x.size, d=hop_size / sample_rate)
    bpms = freqs * 60.0

    mask = (bpms >= 60.0) & (bpms <= 200.0)
    if not np.any(mask):
        return None, 0.0

    filtered = spec[mask]
    filtered_bpms = bpms[mask]
    if filtered.size < 3:
        return None, 0.0

    peak_idx = int(np.argmax(filtered))
    bpm = float(filtered_bpms[peak_idx])

    sorted_peaks = np.sort(filtered)
    second = float(sorted_peaks[-2]) if sorted_peaks.size > 1 else 0.0
    peak = float(filtered[peak_idx])
    confidence = float(np.clip((peak - second) / (peak + 1e-9), 0.0, 1.0))
    return bpm, confidence


def _detect_bpm(y: np.ndarray, sample_rate: int) -> tuple[float | None, float | None]:
    hop_size = 512
    onset = _onset_envelope(y, hop_size=hop_size)
    if onset.size < 16:
        return None, None

    bpm_ac, conf_ac = _tempo_from_autocorrelation(onset, sample_rate=sample_rate, hop_size=hop_size)
    bpm_fft, conf_fft = _tempo_from_frequency(onset, sample_rate=sample_rate, hop_size=hop_size)

    if bpm_ac is None and bpm_fft is None:
        return None, None

    if bpm_ac is not None and bpm_fft is not None:
        diff = abs(bpm_ac - bpm_fft)
        if diff <= 4.0:
            bpm = (bpm_ac + bpm_fft) / 2.0
            agreement = 1.0 - diff / 4.0
            confidence = 0.5 * (conf_ac + conf_fft) * (0.7 + 0.3 * agreement)
        else:
            use_ac = conf_ac >= conf_fft
            bpm = bpm_ac if use_ac else bpm_fft
            confidence = max(conf_ac, conf_fft) * 0.65
    else:
        bpm = bpm_ac if bpm_ac is not None else bpm_fft
        confidence = conf_ac if bpm_ac is not None else conf_fft

    if bpm is None or confidence is None:
        return None, None

    if confidence < 0.22:
        return None, None

    return round(float(bpm), 1), round(float(np.clip(confidence, 0.0, 1.0)), 2)


def _chroma_profile(y: np.ndarray, sample_rate: int, frame_size: int = 4096, hop_size: int = 1024) -> np.ndarray:
    frames = _make_frames(y, frame_size=frame_size, hop_size=hop_size)
    if frames.shape[0] < 3:
        return np.zeros(12, dtype=np.float64)

    window = np.hanning(frame_size).astype(np.float32)
    spectrum = np.abs(np.fft.rfft(frames * window, axis=1)) ** 2
    freqs = np.fft.rfftfreq(frame_size, d=1.0 / sample_rate)

    valid = (freqs >= 40.0) & (freqs <= 5000.0)
    spectrum = spectrum[:, valid]
    freqs = freqs[valid]
    if freqs.size == 0:
        return np.zeros(12, dtype=np.float64)

    midi = np.rint(69.0 + 12.0 * np.log2(freqs / 440.0)).astype(np.int32)
    pitch_classes = np.mod(midi, 12)

    chroma = np.zeros((12, spectrum.shape[0]), dtype=np.float64)
    for pc in range(12):
        pc_bins = pitch_classes == pc
        if not np.any(pc_bins):
            continue
        chroma[pc, :] = spectrum[:, pc_bins].sum(axis=1)

    # Dynamic range compression and time-average.
    chroma = np.sqrt(np.maximum(chroma, 0.0))
    profile = chroma.mean(axis=1)

    total = float(profile.sum())
    if total <= 1e-9:
        return np.zeros(12, dtype=np.float64)

    return profile / total


def _corr(a: np.ndarray, b: np.ndarray) -> float:
    if np.allclose(a.std(), 0.0) or np.allclose(b.std(), 0.0):
        return 0.0
    return float(np.corrcoef(a, b)[0, 1])


def _detect_key(y: np.ndarray, sample_rate: int) -> tuple[str | None, str | None, float | None]:
    profile = _chroma_profile(y, sample_rate=sample_rate)
    if not np.any(profile):
        return None, None, None

    major = _MAJOR_PROFILE / _MAJOR_PROFILE.sum()
    minor = _MINOR_PROFILE / _MINOR_PROFILE.sum()

    candidates: list[tuple[float, str, str]] = []
    for i, key in enumerate(PITCH_CLASSES):
        maj_score = _corr(profile, np.roll(major, i))
        min_score = _corr(profile, np.roll(minor, i))
        candidates.append((maj_score, key, "major"))
        candidates.append((min_score, key, "minor"))

    candidates.sort(key=lambda x: x[0], reverse=True)
    best_score, best_key, best_scale = candidates[0]
    second_score = candidates[1][0] if len(candidates) > 1 else -1.0

    # Confidence blends absolute fit and margin over runner-up.
    margin = max(0.0, best_score - second_score)
    absolute = max(0.0, (best_score + 1.0) / 2.0)
    confidence = float(np.clip(0.45 * absolute + 0.55 * margin, 0.0, 1.0))

    if best_score < 0.18 or confidence < 0.22:
        return None, None, None

    return best_key, best_scale, round(confidence, 2)


def _find_ac_peaks(
    ac: np.ndarray, lag_min: int, lag_max: int, threshold_ratio: float = 0.15
) -> list[tuple[int, float]]:
    """Return (lag, value) pairs for every local maximum above *threshold_ratio* of the segment max."""
    if lag_max >= ac.size:
        lag_max = ac.size - 1
    if lag_max <= lag_min + 1:
        return []

    seg = ac[lag_min : lag_max + 1]
    peak_val = float(np.max(seg))
    if peak_val <= 0:
        return []

    threshold = peak_val * threshold_ratio
    peaks: list[tuple[int, float]] = []
    for i in range(1, seg.size - 1):
        if seg[i] > seg[i - 1] and seg[i] > seg[i + 1] and seg[i] > threshold:
            peaks.append((lag_min + i, float(seg[i])))
    return peaks


def _detect_tala(y: np.ndarray, sample_rate: int) -> dict[str, float | str | None]:
    """Detect rhythm pattern covering both Western meters and Indian talas.

    Returns dict with keys: bpm, bpm_confidence, tala, time_signature.
    """
    hop_size = 512
    onset = _onset_envelope(y, hop_size=hop_size)
    empty: dict[str, float | str | None] = {
        "bpm": None, "bpm_confidence": None, "tala": None, "time_signature": None,
    }
    if onset.size < 32:
        return empty

    x = onset - float(np.mean(onset))
    if np.allclose(x, 0.0, atol=1e-9):
        return empty

    ac = np.correlate(x, x, mode="full")[x.size - 1 :]

    # ---------- Step 1: find the base pulse (akshara / beat) ----------
    # Wide range to cover slow vilambit to fast drut.
    pulse_bpm_min, pulse_bpm_max = 50.0, 280.0
    lag_min = max(1, int(round((60.0 / pulse_bpm_max) * sample_rate / hop_size)))
    lag_max = min(ac.size - 1, int(round((60.0 / pulse_bpm_min) * sample_rate / hop_size)))

    if lag_max <= lag_min:
        return empty

    peaks = _find_ac_peaks(ac, lag_min, lag_max)
    if not peaks:
        return empty

    strongest_score = max(p[1] for p in peaks)
    peaks_by_lag = sorted(peaks, key=lambda p: p[0])

    # Shortest peak that is at least 25 % as strong as the best.
    base_lag = peaks_by_lag[0][0]
    for lag, score in peaks_by_lag:
        if score >= strongest_score * 0.25:
            base_lag = lag
            break

    pulse_bpm = 60.0 * sample_rate / (base_lag * hop_size)
    if pulse_bpm < pulse_bpm_min or pulse_bpm > pulse_bpm_max:
        return empty

    # ---------- Step 2: detect cycle length ----------
    # For each candidate cycle length check if the autocorrelation near
    # (base_lag * cycle_len) is strong relative to the base.
    base_ac = float(ac[base_lag])
    if base_ac <= 0:
        return empty

    best_cycle: int | None = None
    best_ratio = 0.0
    window = max(2, base_lag // 4)

    for cycle_len in sorted(TALA_SIGNATURES.keys()):
        expected_lag = base_lag * cycle_len
        if expected_lag + window >= ac.size:
            continue
        lo = max(0, expected_lag - window)
        hi = min(ac.size, expected_lag + window + 1)
        local_max = float(np.max(ac[lo:hi]))
        # Normalize by base_ac so longer lags (which naturally decay) are
        # not unfairly penalised.  A mild length penalty keeps shorter,
        # simpler cycles preferred when evidence is equal.
        ratio = (local_max / base_ac) * (1.0 - 0.01 * cycle_len)
        if ratio > best_ratio:
            best_ratio = ratio
            best_cycle = cycle_len

    tala: str | None = None
    time_sig: str | None = None
    if best_cycle is not None and best_ratio > 0.15:
        tala, time_sig = TALA_SIGNATURES.get(best_cycle, (None, None))

    # ---------- Step 3: confidence ----------
    seg_mean = float(np.mean(np.abs(ac[lag_min : lag_max + 1])))
    confidence = float(np.clip(
        (base_ac - seg_mean) / (base_ac + 1e-9), 0.0, 1.0,
    ))

    return {
        "bpm": round(pulse_bpm, 1),
        "bpm_confidence": round(confidence, 2),
        "tala": tala,
        "time_signature": time_sig,
    }


def analyze_waveform(y: np.ndarray, sample_rate: int) -> dict[str, float | str | None]:
    if y.size == 0:
        return {
            "bpm": None,
            "bpm_confidence": None,
            "musical_key": None,
            "key_scale": None,
            "key_confidence": None,
        }

    y = np.nan_to_num(y.astype(np.float32, copy=False))
    y = _safe_normalize(y)

    # Avoid intro/outro bias and keep runtime bounded.
    max_samples = int(sample_rate * 180)
    if y.size > max_samples:
        start = (y.size - max_samples) // 2
        y = y[start : start + max_samples]

    # Very short segments are not trustworthy.
    if y.size < sample_rate * 6:
        return {
            "bpm": None,
            "bpm_confidence": None,
            "musical_key": None,
            "key_scale": None,
            "key_confidence": None,
        }

    bpm, bpm_conf = _detect_bpm(y, sample_rate=sample_rate)
    key, scale, key_conf = _detect_key(y, sample_rate=sample_rate)
    tala_result = _detect_tala(y, sample_rate=sample_rate)

    # Use tala detector's BPM as fallback when standard detection fails.
    if bpm is None and tala_result.get("bpm") is not None:
        bpm = tala_result["bpm"]
        bpm_conf = tala_result["bpm_confidence"]

    return {
        "bpm": bpm,
        "bpm_confidence": bpm_conf,
        "musical_key": key,
        "key_scale": scale,
        "key_confidence": key_conf,
        "tala": tala_result.get("tala"),
        "time_signature": tala_result.get("time_signature"),
    }


def analyze_audio_file(audio_path: Path, sample_rate: int = 22050) -> dict[str, float | str | None]:
    y = _decode_audio_mono(audio_path, sample_rate=sample_rate)
    return analyze_waveform(y, sample_rate=sample_rate)