data.py

"""
data.py - workbook loading and the prediction-target registry.

The workbook (AI_Data_Training.xlsx, one row per recording) is read into a tidy table and every
predictable quantity in it is turned into a named target. The registry is data-driven: any
numeric column that is not an identifier becomes a presence (>0) target and a log-level
regression target, and any two-class text column becomes a binary target. Adding a new column
to the workbook (for example a Varroa mite count) therefore adds targets automatically, without
code changes.
"""
import os
import re
import numpy as np
import pandas as pd

MONTHS = ["January", "February", "March", "April", "May", "June", "July", "August",
          "September", "October", "November", "December"]
MONTH_NUM = {m: i + 1 for i, m in enumerate(MONTHS)}
_MONTH_ALIASES = {
    "jan": "January", "feb": "February", "mar": "March", "apr": "April", "may": "May",
    "jun": "June", "june": "June", "jul": "July", "july": "July", "aug": "August",
    "sep": "September", "sept": "September", "oct": "October", "nov": "November",
    "dec": "December",
}
COUNTRY_CODE = {"IA": "USA", "NZ": "NZ"}

ID_COLUMNS = {"colony", "month", "country", "sound_file", "year", "month_num", "t_index",
              "has_audio", "recording_batch"}
KNOWN_BINARY_TEXT = {"y03": "Terminal_or_Survivor_0_to_3_months",
                     "y46": "Terminal_or_Survivor_4_to_6_months"}
FILENAME_FIXES = {  # workbook file name (lower-case) -> file on disk, if ever needed
}


def normalize_month(raw):
    if raw is None or (isinstance(raw, float) and np.isnan(raw)):
        return None
    s = re.sub(r"20\d{2}", "", str(raw)).strip().lower()
    s = re.sub(r"[^a-z]", "", s)
    if s in MONTH_NUM or s.capitalize() in MONTH_NUM:
        return s.capitalize()
    return _MONTH_ALIASES.get(s[:4] if s[:4] in _MONTH_ALIASES else s[:3])


def parse_filename(fname):
    """'12.IA.Sept.2022.wav' -> dict(colony='12', country='USA', month='September', year=2022)."""
    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])
    m = re.search(r"20\d{2}", stem)
    out["year"] = int(m.group()) if m else None
    return out


def _clean_columns(df):
    df = df.copy()
    df.columns = [str(c).strip() for c in df.columns]
    ren = {}
    for c in df.columns:
        lc = c.lower()
        if lc in ("sound file", "sound_file", "file", "filename"):
            ren[c] = "sound_file"
        elif lc == "colony":
            ren[c] = "colony"
        elif lc == "month":
            ren[c] = "month"
        elif lc == "country":
            ren[c] = "country"
        elif c == KNOWN_BINARY_TEXT["y03"]:
            ren[c] = "y03"
        elif c == KNOWN_BINARY_TEXT["y46"]:
            ren[c] = "y46"
    return df.rename(columns=ren)


def load_workbook(xlsx_path, audio_dir=None, sheet=0):
    """Read the workbook into a tidy DataFrame (one row per recording)."""
    df = pd.read_excel(xlsx_path, sheet_name=sheet)
    df = _clean_columns(df)
    if "sound_file" not in df.columns:
        raise ValueError("workbook needs a 'sound file' column")
    df["sound_file"] = df["sound_file"].astype(str).str.strip()
    df["sound_file"] = df["sound_file"].map(lambda s: FILENAME_FIXES.get(s.lower(), s))
    parsed = pd.DataFrame([parse_filename(f) for f in df["sound_file"]])
    if "colony" not in df.columns:
        df["colony"] = parsed["colony"]
    df["colony"] = df["colony"].astype(str).str.strip()
    if "country" not in df.columns:
        df["country"] = parsed["country"]
    df["country"] = df["country"].astype(str).str.strip().str.upper().replace({"IA": "USA", "US": "USA"})
    if "month" not in df.columns:
        df["month"] = parsed["month"]
    df["month"] = df["month"].map(normalize_month)
    df["month_num"] = df["month"].map(MONTH_NUM).astype("Int64")
    df["year"] = parsed["year"].astype("Int64")
    df["t_index"] = ((df["year"] - 2022) * 12 + df["month_num"]).astype("Int64")   # Aug 2022 = 8
    df["recording_batch"] = df["country"] + "_" + df["year"].astype(str) + "_" + df["month_num"].astype(str).str.zfill(2)
    for k in ("y03", "y46"):
        if k in df.columns:
            df[k] = df[k].astype("string").str.strip().str.upper()
    # numeric pathogen / physiology columns: coerce, keep NaN (imputation is a modelling choice)
    for c in df.columns:
        if c in ID_COLUMNS or c in ("y03", "y46"):
            continue
        if df[c].dtype == object:
            conv = pd.to_numeric(df[c], errors="coerce")
            if conv.notna().sum() >= 0.9 * df[c].notna().sum():
                df[c] = conv
    if audio_dir is not None:
        present = {f.lower() for f in os.listdir(audio_dir)}
        df["has_audio"] = df["sound_file"].str.lower().isin(present)
    else:
        df["has_audio"] = True
    return df


# --------------------------------------------------------------------------- targets
def build_targets(df, high_quantile=0.75):
    """Return a list of target dicts: name, kind ('binary'|'regression'), y (float array with
    NaN where undefined), source, description."""
    targets = []
    n = len(df)

    def add(name, kind, y, source, desc):
        y = np.asarray(y, dtype=float)
        targets.append(dict(name=name, kind=kind, y=y, source=source, description=desc))

    for key, col in KNOWN_BINARY_TEXT.items():
        if key in df.columns:
            v = df[key]
            y = np.where(v == "T", 1.0, np.where(v == "S", 0.0, np.nan))
            horizon = "0-3 months" if key == "y03" else "0-6 months (nested: every 0-3 month Terminal row is also Terminal here)"
            add(f"terminal_{key[1:]}", "binary", y, col,
                f"Colony is Terminal (dies) within {horizon} of the recording (1) vs Survivor (0)")
    if "y03" in df.columns and "y46" in df.columns:
        y = np.where((df["y46"] == "T") & (df["y03"] == "S"), 1.0,
                     np.where((df["y46"] == "S") & (df["y03"] == "S"), 0.0, np.nan))
        add("terminal_4_6_only", "binary", y, "Terminal_or_Survivor_4_to_6_months minus 0_to_3",
            "Colony dies 4-6 months after the recording (1) vs survives 6 months (0); rows already Terminal within 3 months excluded")

    numeric_cols = [c for c in df.columns
                    if c not in ID_COLUMNS and c not in ("y03", "y46")
                    and pd.api.types.is_numeric_dtype(df[c])]
    for c in numeric_cols:
        v = pd.to_numeric(df[c], errors="coerce").astype(float).values
        key = re.sub(r"[^a-z0-9]+", "_", c.lower()).strip("_")
        present = np.where(np.isnan(v), np.nan, (v > 0).astype(float))
        add(f"{key}_present", "binary", present, c, f"{c} detected (> 0) in the recording month")
        add(f"{key}_log_level", "regression", np.log1p(np.clip(v, 0, None)), c,
            f"log(1 + {c}) level, all recordings")
        pos = v > 0
        if pos.sum() >= 30:
            thr = float(np.nanquantile(v[pos], high_quantile))
            hi = np.where(pos, (v >= thr).astype(float), np.nan)      # among positives only
            add(f"{key}_high", "binary", hi, c,
                f"Among recordings with {c} > 0: high load (at or above the {int(high_quantile*100)}th percentile of positive values, {thr:.4g}) vs low positive")
            among = np.where(pos, np.log1p(v), np.nan)
            add(f"{key}_log_level_among_positive", "regression", among, c,
                f"log(1 + {c}) among recordings where {c} > 0 (n={int(pos.sum())})")
    if len(numeric_cols) >= 2:
        M = df[numeric_cols].apply(pd.to_numeric, errors="coerce").values.astype(float)
        anyv = np.where(np.isnan(M).all(axis=1), np.nan, (np.nan_to_num(M) > 0).any(axis=1).astype(float))
        add("any_pathogen_present", "binary", anyv, ",".join(numeric_cols),
            "Any of the measured pathogens detected (> 0); largely a site proxy in this data set")

    # other two-class text columns become binary targets automatically
    for c in df.columns:
        if c in ID_COLUMNS or c in ("y03", "y46") or c in numeric_cols:
            continue
        if df[c].dtype == object or str(df[c].dtype) == "string":
            vals = df[c].dropna().astype(str).str.strip().unique()
            if len(vals) == 2:
                pos_label = sorted(vals)[-1]
                y = np.where(df[c].isna(), np.nan, (df[c].astype(str).str.strip() == pos_label).astype(float))
                key = re.sub(r"[^a-z0-9]+", "_", c.lower()).strip("_")
                add(key, "binary", y, c, f"{c} == '{pos_label}' (1) vs '{sorted(vals)[0]}' (0)")

    # confound-detectability probes (not colony-health targets): can audio identify site / season?
    add("probe_country_is_nz", "binary", (df["country"] == "NZ").astype(float).values, "Country",
        "Probe: recording made in NZ (1) vs USA (0); measures site/recording-condition detectability")
    if "recording_batch" in df.columns:
        add("probe_batch_oct2023", "binary", (df["recording_batch"] == "USA_2023_10").astype(float).values, "recording_batch",
            "Probe: recording belongs to the single-visit October 2023 USA batch (1) vs all other recordings (0)")
    if df["month_num"].notna().any():
        mn = df["month_num"].astype(float).values
        add("probe_month_sin", "regression", np.sin(2 * np.pi * mn / 12), "Month",
            "Probe: seasonal phase (sin) of the recording month")
        add("probe_month_cos", "regression", np.cos(2 * np.pi * mn / 12), "Month",
            "Probe: seasonal phase (cos) of the recording month")
    return targets


def metadata_matrix(df):
    """Site, season (two harmonics), year and the October 2023 batch as the metadata-only baseline."""
    mn = df["month_num"].astype(float).fillna(0).values
    ang = 2 * np.pi * mn / 12.0
    yr = df["year"].astype(float).fillna(2022).values
    X = np.column_stack([
        (df["country"] == "NZ").astype(float).values,
        np.sin(ang), np.cos(ang), np.sin(2 * ang), np.cos(2 * ang),
        (yr - 2022.0),
        (df["recording_batch"] == "USA_2023_10").astype(float).values,
    ])
    return X.astype(np.float32), ["is_nz", "month_sin", "month_cos", "month_sin2", "month_cos2", "year_offset", "is_oct2023_batch"]