"""Light feature + sample-analysis engines — CPU, NO torch.

Two modular building blocks for hexa, the same shape/pattern as ears_light:
  • features(path)        → the ~35-dim audio feature stack (/features)
  • sample_profile(path)  → per-sample EDA + measured role (/analyze/samples)

SELF-CONTAINED on purpose: the deployed container holds only armada/api/ (not
armada/tide-table/), so these VENDOR the DSP rather than importing it. They
mirror the canonical tide-table implementations — sample_features.features() and
audio_lens.profile()/classify_family() — and test_feats cross-checks the feature
KEYS against tide-table when it's importable, so drift is caught.

Mastering discipline (feedback_mastering_eda): role/register is decided by the
MEASURED spectrum (centroid + band energy), never inferred from a name. An
optional `name` only disambiguates what spectrum can't (a break is broadband; a
kick and a sub are both sub-band).
"""
import numpy as np

SR = 22050                                   # MIR-standard; plenty for timbre/role
MAX_S = 30.0                                 # cap clip length (loops/long tails)

# Krumhansl-Schmuckler major/minor key profiles (chroma → key estimate)
_KS_MAJ = 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])
_KS_MIN = 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])
_NOTES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

# feature hierarchy (mirrors sample_features.FEATURE_TIERS — robust L0 → derived L1)
FEATURE_TIERS = {
    "L0": ["rms_db", "peak_db", "crest_db", "zcr", "spectral_centroid", "duration_s"],
    "L1": ["log_attack_time", "temporal_centroid", "decay_slope_db_s", "sustain_ratio",
           "spectral_spread", "spectral_skew", "spectral_kurtosis", "spectral_rolloff",
           "spectral_flatness", "spectral_bandwidth", "spectral_flux", "spectral_contrast",
           "chroma_entropy", "key_strength", "f0_median", "hnr", "pct_percussive",
           "onset_rate", "tempo_bpm", "pulse_clarity"] + [f"mfcc_{i}" for i in range(13)],
}

# band edges (Hz) for sample band-energy %, mirrors audio_lens.BANDS
BANDS = [("<80", 0, 80), ("80-150", 80, 150), ("150-500", 150, 500),
         ("500-2k", 500, 2000), ("2k+", 2000, SR / 2)]

# role hints (mirror audio_lens) — identity only disambiguates what spectrum can't
_BREAK_HINTS = ("break", "jungle", "amen", "org_jungle")
_PERC_EXACT = {"kick", "bd", "snare", "sn", "sd", "dr", "drum", "clap", "cp",
               "rim", "rs", "hh", "hat", "hats", "cym", "cy", "tom", "perc",
               "crash", "ride", "shaker", "conga", "bongo"}
_PERC_PREFIX = ("kick", "snare", "clap", "hat", "perc", "rample", "h2ogm",
                "reverbkick", "808kick", "909")
_BASS_HINTS = ("bass", "sub", "808", "reese", "moog", "fbass", "acid", "wobble")


def _load(audio_path, max_seconds, sr=SR):
    import librosa
    y, _ = librosa.load(str(audio_path), sr=sr, mono=True, duration=max_seconds)
    if y.size < sr // 20:                    # need ≥ 50 ms
        raise ValueError("clip too short / undecodable (need ≥ 50 ms of audio)")
    return y


def _env_features(y, sr):
    """Temporal envelope descriptors — the kick↔bass discriminators (Peeters)."""
    import librosa
    rms = librosa.feature.rms(y=y, frame_length=1024, hop_length=256)[0]
    t = librosa.times_like(rms, sr=sr, hop_length=256)
    if rms.sum() < 1e-9 or len(rms) < 3:
        return {}
    pk = int(np.argmax(rms))
    peak = rms[pk] + 1e-9
    attack_t = max(t[pk], 1e-3)
    tempo_cen = float((t * rms).sum() / (rms.sum() + 1e-9))
    tail = rms[pk:]
    if len(tail) > 2:
        tt = t[pk:] - t[pk]
        decay = float(np.polyfit(tt, 20 * np.log10(tail + 1e-9), 1)[0])
    else:
        decay = 0.0
    return {"log_attack_time": float(np.log10(attack_t)),
            "temporal_centroid": tempo_cen,
            "decay_slope_db_s": decay,
            "sustain_ratio": float(np.mean(rms[pk:]) / peak),
            "duration_s": float(len(y) / sr)}


def features(audio_path, max_seconds=MAX_S, rhythm=True) -> dict:
    """Flat dict of ~35 features (+ tiers + engine tag). Robust to short/silent
    clips (a missing feature is omitted, not faked)."""
    import librosa
    y = _load(audio_path, max_seconds)
    f = {}
    # ── time domain ──
    rms = np.sqrt(np.mean(y ** 2)) + 1e-9
    peak = np.abs(y).max() + 1e-9
    f["rms_db"] = float(20 * np.log10(rms))
    f["peak_db"] = float(20 * np.log10(peak))
    f["crest_db"] = float(20 * np.log10(peak / rms))
    f["zcr"] = float(librosa.feature.zero_crossing_rate(y)[0].mean())
    f.update(_env_features(y, SR))
    # ── spectral (moments + shape) ──
    S = np.abs(librosa.stft(y, n_fft=2048, hop_length=512)) ** 2
    freqs = librosa.fft_frequencies(sr=SR, n_fft=2048)
    Sm = S.mean(axis=1)
    tot = Sm.sum() + 1e-12
    cen = float((freqs * Sm).sum() / tot)
    spread = float(np.sqrt(((freqs - cen) ** 2 * Sm).sum() / tot))
    f["spectral_centroid"] = cen
    f["spectral_spread"] = spread
    f["spectral_skew"] = float(((freqs - cen) ** 3 * Sm).sum() / (tot * spread ** 3 + 1e-9))
    f["spectral_kurtosis"] = float(((freqs - cen) ** 4 * Sm).sum() / (tot * spread ** 4 + 1e-9))
    f["spectral_rolloff"] = float(librosa.feature.spectral_rolloff(y=y, sr=SR).mean())
    f["spectral_flatness"] = float(librosa.feature.spectral_flatness(y=y).mean())
    f["spectral_bandwidth"] = float(librosa.feature.spectral_bandwidth(y=y, sr=SR).mean())
    f["spectral_flux"] = float(np.mean(np.diff(S, axis=1).clip(min=0).sum(axis=0)))
    f["spectral_contrast"] = float(librosa.feature.spectral_contrast(y=y, sr=SR).mean())
    # ── cepstral timbre ──
    mfcc = librosa.feature.mfcc(y=y, sr=SR, n_mfcc=13).mean(axis=1)
    for i, v in enumerate(mfcc):
        f[f"mfcc_{i}"] = float(v)
    # ── harmonic / percussive ──
    try:
        yh, yp = librosa.effects.hpss(y)
        eh, ep = float(np.sum(yh ** 2)), float(np.sum(yp ** 2))
        f["pct_percussive"] = float(ep / (eh + ep + 1e-9))
        f["hnr"] = float(eh / (ep + 1e-9))
    except Exception:
        pass
    # ── pitch / key ──
    try:
        chroma = librosa.feature.chroma_cqt(y=y, sr=SR).mean(axis=1)
        cs = chroma / (chroma.sum() + 1e-9)
        f["chroma_entropy"] = float(-(cs * np.log2(cs + 1e-12)).sum())
        cor = [(np.corrcoef(np.roll(_KS_MAJ, k), chroma)[0, 1], _NOTES[k], "maj") for k in range(12)]
        cor += [(np.corrcoef(np.roll(_KS_MIN, k), chroma)[0, 1], _NOTES[k], "min") for k in range(12)]
        best = max(cor, key=lambda x: (x[0] if np.isfinite(x[0]) else -9))
        f["key"], f["key_mode"], f["key_strength"] = best[1], best[2], float(best[0])
    except Exception:
        pass
    try:
        f0 = librosa.yin(y, fmin=40, fmax=2000, sr=SR)
        f0 = f0[np.isfinite(f0)]
        if f0.size:
            f["f0_median"] = float(np.median(f0))
    except Exception:
        pass
    # ── rhythm (loops/takes) — librosa 0.11: feature.tempo, NOT feature.rhythm.tempo ──
    if rhythm:
        try:
            onset_env = librosa.onset.onset_strength(y=y, sr=SR)
            onsets = librosa.onset.onset_detect(onset_envelope=onset_env, sr=SR)
            dur = len(y) / SR
            f["onset_rate"] = float(len(onsets) / dur) if dur else 0.0
            tempo = np.atleast_1d(librosa.feature.tempo(onset_envelope=onset_env, sr=SR))
            f["tempo_bpm"] = float(tempo[0]) if tempo.size else 0.0
            f["pulse_clarity"] = float(np.max(librosa.autocorrelate(onset_env)) /
                                       (np.sum(onset_env ** 2) + 1e-9))
        except Exception:
            pass
    return {"features": f, "tiers": FEATURE_TIERS, "sr": SR, "engine": "light"}


# ── sample analysis (band energy + measured role) ──────────────────────────────
def _profile(y):
    """Band-energy %, spectral centroid, rms/peak — the core 'what is this sound'."""
    w = np.hanning(len(y))
    S = np.abs(np.fft.rfft(y * w)) ** 2
    fr = np.fft.rfftfreq(len(y), 1 / SR)
    tot = S.sum() + 1e-20
    bands = {name: float(100 * S[(fr >= lo) & (fr < hi)].sum() / tot) for name, lo, hi in BANDS}
    return {"rms_db": float(20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9)),
            "peak_db": float(20 * np.log10(np.abs(y).max() + 1e-9)),
            "centroid": float((fr * S).sum() / tot), "bands": bands}


def classify_family(name, prof):
    """(family_key, centroid|None), measurement-aware. Families: percs|bass|melodic|tops|atmos."""
    s = (name or "").lower()
    cen = prof["centroid"] if prof else None
    if any(h in s for h in _BREAK_HINTS):
        return "tops", cen
    if s in _PERC_EXACT or any(s.startswith(h) for h in _PERC_PREFIX):
        return "percs", cen
    bass_name = any(h in s for h in _BASS_HINTS)
    if cen is None:
        return ("bass" if bass_name else "melodic"), None
    if cen < 180:
        return "bass", cen
    if cen < 1400:
        return ("bass" if bass_name else "melodic"), cen
    if cen < 3500:
        return "melodic", cen
    return "atmos", cen


def _time_activity(y, sr):
    """Fraction of the clip above -40 dB of its peak — sustained vs one-shot."""
    import librosa
    rms = librosa.feature.rms(y=y, frame_length=1024, hop_length=256)[0]
    if rms.size == 0:
        return 0.0
    thresh = rms.max() * (10 ** (-40 / 20))
    return float(np.mean(rms > thresh))


def sample_profile(audio_path, name="", max_seconds=MAX_S) -> dict:
    """Per-sample EDA the way mastering needs it: measured spectrum → role, never
    inferred from the name. `name` (optional) only disambiguates breaks/drums."""
    y = _load(audio_path, max_seconds)
    prof = _profile(y)
    family, cen = classify_family(name, prof)
    return {
        "profile": prof,
        "role": family,
        "centroid": cen,
        "broadband": bool(prof["bands"]["150-500"] + prof["bands"]["500-2k"]
                          + prof["bands"]["2k+"] > 18 and prof["centroid"] > 140),
        "time_activity": _time_activity(y, SR),
        "name_hint": name or None,
        "sr": SR,
        "engine": "light",
    }
