bee_features.py

"""
Shared feature extraction + metadata parsing for the Bee colony survival model.

This exact module is used BOTH during training and in the Hugging Face Space so
that audio features are computed identically in both places.

Audio design notes:
  - Honey-bee acoustic activity lives at low frequencies (wing-beat fundamental
    ~180-250 Hz with harmonics extending to ~2 kHz). We therefore resample to
    8 kHz (Nyquist 4 kHz) which preserves all bee-relevant energy while making
    feature extraction ~5x faster than at 44.1 kHz.
  - We aggregate frame-level features over the whole clip (mean/std) to obtain a
    single fixed-length vector per recording, suitable for classical/MLP fusion.
"""
import os
import re
import numpy as np

SR = 8000
N_FFT = 1024
HOP = 512
N_MFCC = 20
N_MELS = 40
MAX_SECONDS = 90            # cap very long clips for speed/consistency
BEE_BANDS_HZ = [(0, 100), (100, 200), (200, 300), (300, 500),
                (500, 1000), (1000, 2000), (2000, 4000)]

# ---- month normalisation ---------------------------------------------------
_MONTH_MAP = {
    'jan': 'January', 'january': 'January',
    'feb': 'February', 'february': 'February',
    'mar': 'March', 'march': 'March',
    'apr': 'April', 'april': 'April',
    'may': 'May',
    'jun': 'June', 'june': 'June',
    'jul': 'July', 'july': 'July',
    'aug': 'August', 'august': 'August',
    'sep': 'September', 'sept': 'September', 'september': 'September',
    'oct': 'October', 'october': 'October',
    'nov': 'November', 'november': 'November',
    'dec': 'December', 'december': 'December',
}
MONTH_ORDER = {m: i for i, m in enumerate(
    ['January', 'February', 'March', 'April', 'May', 'June', 'July',
     'August', 'September', 'October', 'November', 'December'], start=1)}

_COUNTRY_CODE = {'IA': 'USA', 'NZ': 'NZ'}


def normalize_month(raw):
    """Map any month spelling (e.g. 'Sept', 'October2023', 'January ') -> canonical."""
    if raw is None:
        return None
    s = str(raw).strip().lower()
    s = re.sub(r'20\d{2}', '', s).strip()        # drop trailing year e.g. october2023
    s = re.sub(r'[^a-z]', '', s)
    return _MONTH_MAP.get(s, None)


def parse_filename(fname):
    """Parse 'Colony.CountryCode.Month.Year.wav' -> dict of metadata.

    Example: '1.IA.Aug.2022.wav' -> {colony:'1', country:'USA', month:'August', year:2022}
    Returns whatever can be parsed; missing pieces are None.
    """
    base = os.path.basename(str(fname))
    stem = re.sub(r'\.wav$', '', base, flags=re.IGNORECASE)
    parts = stem.split('.')
    out = {'colony': None, 'country': None, 'month': None, 'year': None}
    if len(parts) >= 1:
        out['colony'] = parts[0].strip()
    if len(parts) >= 2:
        out['country'] = _COUNTRY_CODE.get(parts[1].strip().upper(), parts[1].strip())
    if len(parts) >= 3:
        out['month'] = normalize_month(parts[2])
    if len(parts) >= 4:
        m = re.search(r'20\d{2}', parts[3])
        out['year'] = int(m.group()) if m else None
    # year sometimes appended to the month token (october2023)
    if out['year'] is None:
        m = re.search(r'20\d{2}', stem)
        if m:
            out['year'] = int(m.group())
    return out


# ---- audio feature extraction ---------------------------------------------
def _agg(name, arr):
    """mean/std of a (n_features, n_frames) or (n_frames,) array -> flat dict."""
    arr = np.atleast_2d(arr)
    out = {}
    mean = np.nanmean(arr, axis=1)
    std = np.nanstd(arr, axis=1)
    for i in range(arr.shape[0]):
        suffix = f'{i}' if arr.shape[0] > 1 else ''
        out[f'{name}{suffix}_mean'] = float(mean[i])
        out[f'{name}{suffix}_std'] = float(std[i])
    return out


def extract_features_from_audio(y, sr):
    """Compute the fixed-length bee feature dict from a waveform."""
    import librosa
    if sr != SR:
        y = librosa.resample(y, orig_sr=sr, target_sr=SR)
        sr = SR
    if y.ndim > 1:                      # to mono
        y = np.mean(y, axis=0) if y.shape[0] < y.shape[1] else np.mean(y, axis=1)
    y = y.astype(np.float32)
    if MAX_SECONDS is not None:
        y = y[: MAX_SECONDS * sr]
    # normalise amplitude so loudness/mic-gain differences don't dominate
    peak = np.max(np.abs(y)) if y.size else 0.0
    if peak > 0:
        y = y / peak

    feats = {}
    S = np.abs(librosa.stft(y, n_fft=N_FFT, hop_length=HOP)) ** 2  # power spectrogram
    mel = librosa.feature.melspectrogram(S=S, sr=sr, n_mels=N_MELS)
    logmel = librosa.power_to_db(mel + 1e-10)
    mfcc = librosa.feature.mfcc(S=librosa.power_to_db(mel + 1e-10), n_mfcc=N_MFCC)
    dmfcc = librosa.feature.delta(mfcc)

    feats.update(_agg('mfcc', mfcc))
    feats.update(_agg('dmfcc', dmfcc))
    feats.update(_agg('logmel', logmel))
    feats.update(_agg('centroid', librosa.feature.spectral_centroid(S=np.sqrt(S), sr=sr)))
    feats.update(_agg('bandwidth', librosa.feature.spectral_bandwidth(S=np.sqrt(S), sr=sr)))
    feats.update(_agg('rolloff', librosa.feature.spectral_rolloff(S=np.sqrt(S), sr=sr)))
    feats.update(_agg('flatness', librosa.feature.spectral_flatness(S=np.sqrt(S))))
    feats.update(_agg('contrast', librosa.feature.spectral_contrast(
        S=np.sqrt(S), sr=sr, fmin=100.0, n_bands=4)))
    feats.update(_agg('zcr', librosa.feature.zero_crossing_rate(y, frame_length=N_FFT, hop_length=HOP)))
    feats.update(_agg('rms', librosa.feature.rms(S=np.sqrt(S), frame_length=N_FFT, hop_length=HOP)))

    # bee-band relative energies from the mean power spectrum
    freqs = librosa.fft_frequencies(sr=sr, n_fft=N_FFT)
    mean_spec = np.mean(S, axis=1)
    total = mean_spec.sum() + 1e-12
    for lo, hi in BEE_BANDS_HZ:
        band = mean_spec[(freqs >= lo) & (freqs < hi)].sum()
        feats[f'bandfrac_{lo}_{hi}'] = float(band / total)
    feats['dominant_freq'] = float(freqs[int(np.argmax(mean_spec))])
    feats['spectral_entropy'] = float(
        -np.sum((mean_spec / total) * np.log(mean_spec / total + 1e-12)))
    return feats


def extract_features_from_file(path):
    """Load an audio file and return its feature dict (mono, resampled)."""
    import librosa
    y, sr = librosa.load(path, sr=SR, mono=True)
    return extract_features_from_audio(y, sr)


# canonical, sorted feature-name list so training and inference agree on order
def audio_feature_names():
    dummy = np.zeros(SR * 2, dtype=np.float32)
    return sorted(extract_features_from_audio(dummy, SR).keys())


# ---- unified model-input construction (shared by training & deployment) -----
# Tabular block, in fixed order. Audio block is appended in audio_feature_names()
# order. build_feature_vector() below is the SINGLE source of truth for the model
# input, guaranteeing identical preprocessing at train and inference time.
TABULAR_ORDER = ['is_usa', 'is_nz', 'month_sin', 'month_cos',
                 'log_cbpv', 'log_dwv', 'log_kbv']


def build_tabular_dict(country, month, cbpv, dwv, kbv):
    """Turn raw metadata into the numeric tabular feature dict."""
    country = (str(country).strip().upper() if country is not None else '')
    is_usa = 1.0 if country in ('USA', 'US', 'IA') else 0.0
    is_nz = 1.0 if country in ('NZ', 'NEW ZEALAND') else 0.0
    mnum = MONTH_ORDER.get(normalize_month(month) or month, 0)
    ang = 2.0 * np.pi * (mnum / 12.0)
    def _log(x):
        try:
            return float(np.log1p(max(0.0, float(x))))
        except (TypeError, ValueError):
            return 0.0
    return {
        'is_usa': is_usa, 'is_nz': is_nz,
        'month_sin': float(np.sin(ang)) if mnum else 0.0,
        'month_cos': float(np.cos(ang)) if mnum else 0.0,
        'log_cbpv': _log(cbpv), 'log_dwv': _log(dwv), 'log_kbv': _log(kbv),
    }


def build_feature_vector(tabular_dict, audio_dict, audio_names, use_audio=True):
    """Concatenate tabular + audio into a single ordered numpy vector."""
    vec = [tabular_dict.get(k, 0.0) for k in TABULAR_ORDER]
    if use_audio:
        vec += [audio_dict.get(k, 0.0) for k in audio_names]
    return np.asarray(vec, dtype=np.float32)


def full_feature_order(audio_names, use_audio=True):
    return list(TABULAR_ORDER) + (list(audio_names) if use_audio else [])