"""roles — map a producer-labelled stem filename to a role family, then VERIFY it.

Fred again's dropbox (and any producer stem pack) ships stems named by the human
who made them: `MAREA MIX10 123BPM KIT STEM.wav`, `ANGIE MIX18 PAD STUFF STEM.wav`.
That label is a strong prior and a terrible fact. The house rule stands — never
infer a sound's role from its name, validate by analysis (`feedback_mastering_eda`)
— so this module does both halves:

  1. `parse_stem_name()`  — filename → (track, subrole token) by stripping the
     `<TRACK> MIX<n> [<bpm>BPM] <ROLE> STEM` scaffolding the pack uses.
  2. `ROLE_MAP`           — subrole token → the *claimed* family + a mode hint.
  3. `probe()`            — cheap EDA (level, silence, centroid, LF fraction,
     onset density) that CONFIRMS or CONTRADICTS the claim, and is the only thing
     downstream code is allowed to trust.

Families are the finder's four canonical stems plus `fx`, because a drone/noise/
texture stem behaves like none of the four and wants its own kit.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Optional

import numpy as np

# ── families ─────────────────────────────────────────────────────────────────
# `stem` is what the finder/export code keys on (bass mono-sum, rotation priority,
# grid preference); `family` is what names the kit.
FAMILIES = {
    "drums": {"stem": "drums", "mode": "loops"},
    "bass":  {"stem": "bass",  "mode": "loops"},
    "vox":   {"stem": "vocals", "mode": "both"},
    "tonal": {"stem": "other", "mode": "loops"},
    "fx":    {"stem": "other", "mode": "both"},
    "hits":  {"stem": "drums", "mode": "chops"},   # isolated one-hit stems
}

# subrole token → (family, is_composite)
# `is_composite` marks a stem that already contains a full groove/arrangement —
# those are the loop gold (ALL DRUMS, KIT, BREAKS) vs an isolated element (KICK).
ROLE_MAP: dict[str, tuple[str, bool]] = {
    # rhythmic composites — the loop gold
    "all drums": ("drums", True), "drums": ("drums", True), "kit": ("drums", True),
    "drum breaks": ("drums", True), "breaks": ("drums", True),
    "main drum loop": ("drums", True), "drum loop 1": ("drums", True),
    "lo drums": ("drums", True), "hi drums": ("drums", True),
    "perc": ("drums", True), "hihats etc": ("drums", True), "hihats": ("drums", True),
    # isolated hits — one-shot material, not loops
    "kick": ("hits", False), "kicks": ("hits", False), "snare": ("hits", False),
    "clap": ("hits", False), "toms": ("hits", False),
    # low end
    "bass": ("bass", True),
    # voice
    "lv": ("vox", True), "bv": ("vox", True), "f vox": ("vox", True),
    "voices": ("vox", True), "choir": ("vox", True), "breaths": ("vox", False),
    "vox textures": ("vox", True), "vox": ("vox", True),
    # tonal
    "synth": ("tonal", True), "keys": ("tonal", True), "keys1": ("tonal", True),
    "keys2": ("tonal", True), "pad": ("tonal", True), "pad stuff": ("tonal", True),
    "strings": ("tonal", True), "guitar": ("tonal", True), "organ": ("tonal", True),
    "piano": ("tonal", True),
    # texture / effect
    "fx": ("fx", False), "drone": ("fx", False), "drones": ("fx", False),
    "noise": ("fx", False), "textures": ("fx", True), "ipad textures": ("fx", True),
    "morph": ("fx", False),
}

# scaffolding this pack wraps every role in
_SCAFFOLD = re.compile(
    r"""^\s*
    (?P<track>.*?)                 # TRACK (greedy-least)
    \s+(?:MIX\d*|NEW)\s+           # MIX18 | MIX1 | NEW
    (?:\d+\s*BPM\s+)?              # optional 123BPM
    (?P<role>.+?)                  # the role words
    \s+STEM\s*$""",
    re.IGNORECASE | re.VERBOSE,
)
_BPM_IN_NAME = re.compile(r"(\d{2,3})\s*BPM", re.IGNORECASE)


@dataclass
class StemName:
    path_stem: str
    track: Optional[str]
    role_token: str            # normalised, lowercase ("all drums")
    claimed_family: str        # from ROLE_MAP, or "tonal" when unknown
    composite: bool
    known: bool                # was the token in ROLE_MAP?
    bpm_hint: Optional[float]  # producer-declared BPM, if the filename carried one


def parse_stem_name(filename: str) -> StemName:
    """`MAREA MIX10 123BPM KIT STEM.wav` → track=MAREA role=kit bpm_hint=123."""
    base = re.sub(r"\.wav$", "", filename, flags=re.IGNORECASE)
    bpm = _BPM_IN_NAME.search(base)
    m = _SCAFFOLD.match(base)
    if m:
        track, role = m.group("track"), m.group("role")
    else:                                  # unscaffolded: last resort, whole name
        track, role = None, re.sub(r"\s+STEM$", "", base, flags=re.IGNORECASE)
    token = re.sub(r"\s+", " ", role.strip().lower())
    fam, comp = ROLE_MAP.get(token, (None, None))
    # a track-specific token (MAREA STEM, TONI VOX, MAKE IT THRU VOX) — fall back on
    # the words it contains rather than giving up.
    if fam is None:
        for key, (f, c) in ROLE_MAP.items():
            if re.search(rf"\b{re.escape(key)}\b", token):
                fam, comp = f, c
                break
    return StemName(
        path_stem=base, track=track, role_token=token,
        claimed_family=fam or "tonal", composite=True if comp is None else comp,
        known=fam is not None,
        bpm_hint=float(bpm.group(1)) if bpm else None,
    )


# ── the verification half ────────────────────────────────────────────────────
@dataclass
class StemProbe:
    """Cheap EDA that either backs the filename's claim or overrides it."""
    rms_dbfs: float
    peak_dbfs: float
    active_frac: float        # share of frames within 40 dB of this stem's own peak
    active_rms_dbfs: float    # how loud it is WHEN it plays
    centroid_hz: float
    lf_frac: float            # energy below 150 Hz / total
    hf_frac: float            # energy above 2 kHz / total
    onset_rate: float         # onsets per second over the loud excerpts
    env_crest: float          # frame-RMS crest: ~1 sustained, >3 struck (gates onset_rate)
    measured_family: str
    agrees: bool = field(default=False)
    note: str = ""

    @property
    def usable(self) -> bool:
        """Is there material here at all?

        Judged on PEAK and on how loud the stem is when it plays — not on overall RMS,
        which on a sparse stem measures how much silence it contains. ANGIE's HIHATS
        stem is −48.8 dBFS RMS and 9.7 % active, and it is a perfectly good hi-hat
        stem peaking at −14.7 (`feedback_right_lens_per_control`).
        """
        return self.peak_dbfs > -40.0 and self.active_frac >= 0.005


def _loud_windows(db: np.ndarray, y: np.ndarray, sr: int, hop: int, *,
                  n: int = 4, win_s: float = 20.0) -> np.ndarray:
    """Excerpts chosen where the stem is PLAYING, concatenated.

    Classifying a five-minute stem does not need five minutes of spectra, but the
    cheap version — evenly spaced windows — samples the gaps. Measured on ANGIE:
    even windows read HIHATS at −60.9 dBFS / 1.4 % active where the full file says
    −48.8 / 9.7 %, and missed PAD STUFF's material completely (0.0 % vs 6.9 %),
    because that stem only plays in two sections. A stem is not stationary, so where
    you look decides what you find (`feedback_measure_the_time_axis`).

    So: rank non-overlapping windows by energy and take the loudest few.
    """
    win = int(win_s * sr)
    if len(y) <= win * n:
        return y
    per = max(1, win // hop)
    nwin = len(db) // per
    if nwin <= n:
        return y
    energy = np.array([db[i * per:(i + 1) * per].mean() for i in range(nwin)])
    picks = sorted(np.argsort(energy)[-n:])
    return np.concatenate([y[i * per * hop:i * per * hop + win] for i in picks])


def probe(y: np.ndarray, sr: int, claimed: str) -> StemProbe:
    """Measure a mono stem and decide what it actually is. Name is never trusted.

    Note that `env_crest` is a sparsity measure, not a percussion one — MAREA's single
    vocal phrase reads 17.98 — which is why it is only ever used to *gate* onset density,
    never on its own.
    """
    import librosa

    eps = 1e-12
    hop = 512
    peak = float(20 * np.log10(np.max(np.abs(y)) + eps))
    rms = float(20 * np.log10(np.sqrt(np.mean(y ** 2)) + eps))

    # Level and activity come from the WHOLE file — it is only a frame-RMS pass, and
    # it is the number that decides whether a stem is worth cutting at all.
    frame_db = 20 * np.log10(librosa.feature.rms(y=y, hop_length=hop)[0] + eps)
    # "Active" relative to this stem's own peak, not an absolute floor: an isolated
    # hi-hat stem is mostly silence by design, and its overall RMS is a statement about
    # how much silence it contains, not about whether it has material in it.
    active = (frame_db > peak - 40.0) & (frame_db > -60.0)
    active_frac = float(active.mean()) if active.size else 0.0
    active_rms = float(frame_db[active].mean()) if active.any() else -120.0

    # Spectral shape and onsets, on the loud excerpts only.
    w = _loud_windows(frame_db, y, sr, hop)
    S = np.abs(librosa.stft(w, n_fft=2048, hop_length=hop)) ** 2
    freqs = librosa.fft_frequencies(sr=sr, n_fft=2048)
    tot = S.sum() + eps
    band = S.sum(axis=1)
    centroid = float((band * freqs).sum() / tot)
    lf_frac = float(band[freqs < 150].sum() / tot)
    hf_frac = float(band[freqs > 2000].sum() / tot)

    onsets = librosa.onset.onset_detect(y=w, sr=sr, units="time", backtrack=True)
    onset_rate = float(len(onsets) / max(len(w) / sr, 1.0))

    # Onset COUNT is not a percussiveness measure. librosa's peak-picker finds peaks in
    # whatever envelope it is handed, including a flat one: a steady 50 Hz sine yields 84
    # onsets over six seconds (14/s), which would file a pure bass tone as drums. So the
    # count is gated on the envelope actually having transients in it — the crest of the
    # frame-RMS, ~1.0 for a sustained tone against 3.1-9.5 for anything struck.
    fr = librosa.feature.rms(y=w, hop_length=hop)[0]
    env_crest = float(np.percentile(fr, 99) / (fr.mean() + eps))

    # ── the verdict ──────────────────────────────────────────────────────────
    # PERCUSSIVE FIRST. The obvious test — "mostly low energy, low centroid ⇒ bass" —
    # fires on a drum bus, because a modern kick carries most of a drum mix's ENERGY:
    # measured on MAREA and ANGIE, ALL DRUMS reads lf_frac 0.84/0.86 and centroid
    # 182/240 Hz, indistinguishable from a bass stem on those two features alone. What
    # separates them is onset density (ALL DRUMS 6.6/s, BASS 1.2/s), so density is
    # asked first and the spectral test only sees what is left.
    if onset_rate > 3.0 and env_crest > 2.0:
        measured = "drums"
    elif hf_frac > 0.25 and onset_rate > 0.8:
        measured = "drums"           # bright and struck: a clap/snare stem is sparse
    elif lf_frac > 0.75 and hf_frac < 0.05:
        measured = "bass"
    elif onset_rate < 0.35:
        measured = "fx"              # sustained/sparse: pad, drone, texture
    elif onset_rate > 1.2 and centroid > 1200:
        measured = "drums"
    else:
        measured = "tonal"

    p = StemProbe(rms_dbfs=round(rms, 2), peak_dbfs=round(peak, 2),
                  active_frac=round(active_frac, 3), active_rms_dbfs=round(active_rms, 2),
                  centroid_hz=round(centroid, 1), lf_frac=round(lf_frac, 3),
                  hf_frac=round(hf_frac, 4), onset_rate=round(onset_rate, 2),
                  env_crest=round(env_crest, 2), measured_family=measured)

    # `vox` and `tonal` are not separable by these cheap features (a sung line and a
    # rhodes chord have the same coarse shape) — CLAP does that job downstream. So a
    # vox claim is only CONTRADICTED when the measurement says bass or drums.
    claim_stem = FAMILIES.get(claimed, FAMILIES["tonal"])["stem"]
    meas_stem = FAMILIES.get(measured, FAMILIES["tonal"])["stem"]
    soft = {"vox", "tonal", "fx"}
    p.agrees = (claim_stem == meas_stem) or (claimed in soft and measured in soft)
    if not p.agrees:
        p.note = f"filename says {claimed}, measurement says {measured}"
    if not p.usable:
        p.note = (p.note + "; " if p.note else "") + \
            f"no material (peak {peak:.0f} dBFS, active {active_frac:.1%})"
    return p


def effective_family(name: StemName, p: StemProbe) -> str:
    """What we actually treat the stem as — and how much authority measurement has.

    The house rule is that a sound's role is never *inferred* from its name; it is not
    that a producer's label about their own stem is worthless. These two claims have
    very different evidential weight, and the split matters:

    * **Presence is measured, always.** Whether there is material in a stem is decided
      by the probe and nothing else.
    * **Bass is measured.** "Almost all energy under 150 Hz, essentially nothing above
      2 kHz, not densely struck" is unambiguous, and a stem like that is bass whatever
      the filename says.
    * **Everything else keeps the label.** Onset density cannot separate a rhythmic
      rhodes from a drum bus — measured on MAREA, KEYS2 reads 5.1 onsets/s and would be
      filed as drums, and the sparse CLAP stem reads 1.0/s and would be filed as tonal.
      Both are wrong, and both would be wrong *silently*, which is worse than the
      label's known imprecision.

    Every disagreement is reported either way, so a wrong label stays visible instead of
    being quietly overwritten by a wrong measurement.
    """
    if p.agrees:
        return name.claimed_family
    if not name.known:
        # not the producer's word — our own fallback for a token like `MAREA STEM`.
        # A guess loses to a measurement.
        return p.measured_family
    if (name.claimed_family in ("tonal", "fx") and p.measured_family == "bass"
            and p.lf_frac > 0.75 and p.hf_frac < 0.05):
        return "bass"
    return name.claimed_family
