"""
europa.analysis
~~~~~~~~~~~~~~~
High-level audio analysis.

Functions:
    beat_track      - Estimate BPM and beat locations
    onset_detect    - Detect note/event onsets
    pitch_track     - Track fundamental frequency (f0) over time
    key_detect      - Estimate musical key
    tempo           - Estimate global tempo in BPM
    silence_frames  - Detect silent frames
    detect_clipping - Find clipped samples
"""

import numpy as np
from ..core.audio import Audio
from ..utils.dsp import stft as _stft


def beat_track(
    audio: Audio,
    hop_length: int = 512,
    start_bpm: float = 120.0,
) -> tuple:
    """
    Estimate tempo and beat locations.

    Args:
        audio      : Input Audio.
        hop_length : Hop size for onset detection.
        start_bpm  : Initial BPM guess.

    Returns:
        tuple: (bpm: float, beat_times: np.ndarray)
               beat_times is an array of beat positions in seconds.

    Examples:
        >>> bpm, beats = europa.analysis.beat_track(audio)
        >>> print(f"Tempo: {bpm:.1f} BPM")
    """
    onset_env = _onset_envelope(audio, hop_length=hop_length)
    bpm = tempo(audio, hop_length=hop_length, start_bpm=start_bpm)
    period = int(60.0 / bpm * audio.sr / hop_length)

    # Dynamic programming beat tracker
    beats = _dp_beat_track(onset_env, period)
    beat_times = beats * hop_length / audio.sr
    return bpm, beat_times


def tempo(
    audio: Audio,
    hop_length: int = 512,
    start_bpm: float = 120.0,
) -> float:
    """
    Estimate global tempo in BPM.

    Args:
        audio      : Input Audio.
        hop_length : Hop size.
        start_bpm  : Initial BPM estimate for autocorrelation.

    Returns:
        float: Estimated tempo in BPM.
    """
    onset_env = _onset_envelope(audio, hop_length=hop_length)
    # Autocorrelation-based tempo estimation
    ac = _autocorrelate(onset_env)
    bpm_range = np.array([30.0, 300.0])
    lag_range = (60.0 * audio.sr / hop_length / bpm_range).astype(int)
    lag_range = np.clip(lag_range, 1, len(ac) - 1)

    valid = ac[lag_range[1]:lag_range[0]]
    best_lag = np.argmax(valid) + lag_range[1]
    bpm = 60.0 * audio.sr / (hop_length * best_lag)
    return float(bpm)


def onset_detect(
    audio: Audio,
    hop_length: int = 512,
    threshold: float = 0.5,
    wait: int = 8,
) -> np.ndarray:
    """
    Detect onset times (note/event start positions).

    Args:
        audio      : Input Audio.
        hop_length : Hop size.
        threshold  : Peak-picking threshold (0.0–1.0).
        wait       : Minimum frames between consecutive onsets.

    Returns:
        np.ndarray: Onset times in seconds.
    """
    onset_env = _onset_envelope(audio, hop_length=hop_length)
    # Normalize
    onset_env /= (np.max(onset_env) + 1e-9)
    # Peak picking
    peaks = _peak_pick(onset_env, threshold=threshold, wait=wait)
    return peaks * hop_length / audio.sr


def pitch_track(
    audio: Audio,
    hop_length: int = 512,
    fmin: float = 80.0,
    fmax: float = 2000.0,
) -> tuple:
    """
    Track fundamental frequency (f0) using autocorrelation.

    Args:
        audio      : Input Audio (mono preferred).
        hop_length : Hop size between frames.
        fmin       : Minimum expected pitch in Hz.
        fmax       : Maximum expected pitch in Hz.

    Returns:
        tuple: (times: np.ndarray, frequencies: np.ndarray)
               Unvoiced frames are set to 0.0 Hz.
    """
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    frame_length = 2048

    lag_min = int(audio.sr / fmax)
    lag_max = int(audio.sr / fmin)

    n_frames = 1 + (len(data) - frame_length) // hop_length
    freqs = np.zeros(n_frames)
    times = np.arange(n_frames) * hop_length / audio.sr

    for i in range(n_frames):
        start = i * hop_length
        frame = data[start:start + frame_length]
        if len(frame) < frame_length:
            break
        ac = _autocorrelate(frame)
        if len(ac) <= lag_max:
            continue
        ac_search = ac[lag_min:lag_max]
        best_lag = np.argmax(ac_search) + lag_min
        # Voiced/unvoiced detection
        if ac[best_lag] > 0.3 * ac[0]:
            freqs[i] = audio.sr / best_lag

    return times, freqs


def key_detect(audio: Audio) -> str:
    """
    Estimate the musical key of the audio using chromagram analysis.

    Returns:
        str: Estimated key (e.g. 'C major', 'A minor').
    """
    from ..feature import chroma as _chroma
    chroma = _chroma(audio)
    chroma_mean = np.mean(chroma, axis=1)

    # 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])
    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])

    note_names = ['C','C#','D','D#','E','F','F#','G','G#','A','A#','B']
    best_score = -np.inf
    best_key   = "C major"

    for i in range(12):
        maj = np.roll(major_profile, i)
        min_ = np.roll(minor_profile, i)
        score_maj = np.corrcoef(chroma_mean, maj)[0, 1]
        score_min = np.corrcoef(chroma_mean, min_)[0, 1]
        if score_maj > best_score:
            best_score = score_maj
            best_key   = f"{note_names[i]} major"
        if score_min > best_score:
            best_score = score_min
            best_key   = f"{note_names[i]} minor"

    return best_key


def silence_frames(
    audio: Audio,
    hop_length: int = 512,
    top_db: float = 60.0,
) -> np.ndarray:
    """
    Return a boolean mask of silent frames.

    Returns:
        np.ndarray: Boolean array of shape (n_frames,). True = silent.
    """
    from ..feature import rms_energy
    rms = rms_energy(audio, hop_length=hop_length)[0]
    rms_db = 20 * np.log10(rms + 1e-10)
    ref_db = np.max(rms_db)
    return rms_db < (ref_db - top_db)


def detect_clipping(audio: Audio, threshold: float = 0.99) -> np.ndarray:
    """
    Detect clipped (saturated) samples.

    Args:
        audio     : Input Audio.
        threshold : Amplitude threshold above which samples are clipped.

    Returns:
        np.ndarray: Indices of clipped samples.
    """
    return np.where(np.abs(audio.data) >= threshold)[0]


# ── Internal helpers ──────────────────────────────────────────────────────────

def _onset_envelope(audio: Audio, hop_length: int = 512) -> np.ndarray:
    """Compute a spectral flux onset envelope."""
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    D = np.abs(_stft(data, hop_length=hop_length))
    # Spectral flux: positive differences in magnitude spectrum
    diff = np.diff(D, axis=1)
    flux = np.sum(np.maximum(diff, 0), axis=0)
    return flux


def _autocorrelate(x: np.ndarray) -> np.ndarray:
    """Normalized autocorrelation."""
    n = len(x)
    result = np.correlate(x, x, mode='full')
    result = result[n - 1:]
    result /= (result[0] + 1e-9)
    return result


def _peak_pick(x: np.ndarray, threshold: float = 0.5, wait: int = 8) -> np.ndarray:
    """Simple peak picking algorithm."""
    peaks = []
    last = 0
    for i in range(1, len(x) - 1):
        if x[i] > threshold and x[i] >= x[i - 1] and x[i] >= x[i + 1]:
            if i - last >= wait:
                peaks.append(i)
                last = i
    return np.array(peaks)


def _dp_beat_track(onset_env: np.ndarray, period: int) -> np.ndarray:
    """Simple dynamic programming beat tracker."""
    n = len(onset_env)
    if n == 0:
        return np.array([])

    score = onset_env.copy()
    prev  = np.arange(n)

    for i in range(1, n):
        best_prev = max(0, i - int(period * 1.5))
        j = np.argmax(score[best_prev:i]) + best_prev
        score[i] += score[j]
        prev[i] = j

    # Backtrack
    beats = []
    i = np.argmax(score)
    while i > 0 and prev[i] != i:
        beats.append(i)
        i = prev[i]
    beats.append(i)
    return np.array(sorted(beats))
