"""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        # fraction of frames above -50 dBFS
    centroid_hz: float
    lf_frac: float            # energy below 150 Hz / total
    onset_rate: float         # onsets per second over ACTIVE time
    measured_family: str
    agrees: bool = field(default=False)
    note: str = ""

    @property
    def usable(self) -> bool:
        return self.active_frac >= 0.05 and self.rms_dbfs > -45.0


def _probe_windows(y: np.ndarray, sr: int, *, n: int = 4, win_s: float = 20.0
                   ) -> np.ndarray:
    """Evenly-spaced excerpts, concatenated — the probe's view of a long stem.

    Classifying a 5-minute stem does not need 5 minutes of spectra; it needs a fair
    sample of them. Several spread windows rather than one long middle slice, because
    a stem's character is not stationary — a pad that only enters at the drop reads as
    silence if you look at one place (`feedback_measure_the_time_axis`). Short stems
    fall through untouched.
    """
    win = int(win_s * sr)
    if len(y) <= win * n:
        return y
    starts = np.linspace(0, len(y) - win, n).astype(int)
    return np.concatenate([y[s:s + win] for s in starts])


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

    y = _probe_windows(y, sr)
    eps = 1e-12
    peak = 20 * np.log10(np.max(np.abs(y)) + eps)
    rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + eps)

    hop = 512
    frame_rms = librosa.feature.rms(y=y, hop_length=hop)[0]
    frame_db = 20 * np.log10(frame_rms + eps)
    active = frame_db > -50.0
    active_frac = float(active.mean()) if active.size else 0.0

    # spectral shape measured on ACTIVE frames only — a stem that is 80% silence
    # (an FX hit, a one-bar riser) would otherwise report the centroid of its noise
    # floor (`feedback_measure_the_time_axis`).
    S = np.abs(librosa.stft(y, n_fft=2048, hop_length=hop))
    if active.size and active.any():
        S = S[:, : active.size][:, active[: S.shape[1]]] if S.shape[1] >= active.size \
            else S[:, active[: S.shape[1]]]
    if S.size == 0:
        S = np.abs(librosa.stft(y, n_fft=2048, hop_length=hop))
    freqs = librosa.fft_frequencies(sr=sr, n_fft=2048)
    power = S ** 2
    tot = power.sum() + eps
    centroid = float((power.sum(axis=1) * freqs).sum() / tot)
    lf_frac = float(power[freqs < 150].sum() / tot)

    onsets = librosa.onset.onset_detect(y=y, sr=sr, units="time", backtrack=True)
    active_s = max(active_frac * len(y) / sr, 1e-3)
    onset_rate = float(len(onsets) / active_s)

    # ── the verdict ──────────────────────────────────────────────────────────
    # Ordered most-specific first. These thresholds are deliberately loose: the
    # probe's job is to catch a stem that is nothing like its label, not to
    # re-derive a taxonomy the producer already knows.
    if lf_frac > 0.80 and centroid < 250:
        measured = "bass"
    elif onset_rate > 1.2 and centroid > 1200:
        measured = "drums"
    elif onset_rate > 1.2 and lf_frac > 0.5:
        measured = "drums"           # kick-led groove: percussive but dark
    elif onset_rate < 0.35:
        measured = "fx"              # sustained/sparse: pad, drone, texture
    else:
        measured = "tonal"

    p = StemProbe(rms_dbfs=round(rms, 2), peak_dbfs=round(peak, 2),
                  active_frac=round(active_frac, 3), centroid_hz=round(centroid, 1),
                  lf_frac=round(lf_frac, 3), onset_rate=round(onset_rate, 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"near-silent (active {active_frac:.0%}, rms {rms:.0f} dBFS)"
    return p


def effective_family(name: StemName, p: StemProbe) -> str:
    """What we actually treat the stem as. Measurement wins on a hard contradiction.

    `hits` is a *name-only* distinction (a KICK stem and an ALL DRUMS stem measure
    almost identically — both are percussive and dark) so it survives a `drums`
    measurement; that is the one place the label carries information the cheap
    features cannot.
    """
    if p.agrees:
        return name.claimed_family
    if name.claimed_family == "hits" and p.measured_family == "drums":
        return "hits"
    return p.measured_family
