features.py

"""Derive the Phase-1 linguistic factors from an audio file + its transcript.

No transcription happens here: the student uploads the existing word-timestamped
transcript (JSON), and we run markers.extract_markers (disfluency, fluency,
complexity + Phase-1 acoustic). No MFA / forced alignment, so Phase-2 markers
(VOT, vowel space, rhythm) are not produced.

The output is one flat dict per recording, ready to stack into a table for
modeling. Column names are kept short and spreadsheet-friendly.
"""
from __future__ import annotations

import json
import os

import librosa

from markers import extract_markers

# (column label, dotted path into extract_markers() output, decimals)
FLATTEN = [
    ("Words", "word_count_clean", 0),
    ("Duration_s", "duration_s", 1),
    # disfluency (per 100 words)
    ("FP_per_100w", "filled_pauses.per_100_words", 2),
    ("EP_per_100w", "empty_pauses.per_100_words", 2),
    ("RP_per_100w", "repetitions.per_100_words", 2),
    ("RT_per_100w", "retractions.per_100_words", 2),
    ("PauseMean_ms", "temporal.pause_mean_ms", 1),
    # fluency
    ("SpeechRate_spm", "temporal.speech_rate_spm", 2),
    ("ArticRate_spm", "temporal.articulation_rate_spm", 2),
    ("Phonation_pct", "temporal.phonation_time_ratio_pct", 2),
    # complexity
    ("MLU", "temporal.mlu_words", 2),
    ("TTR", "lexical_diversity.ttr", 4),
    ("MTLD", "lexical_diversity.mtld", 2),
    ("CodeSwitch", "code_switching.count", 0),
    # acoustic: pitch
    ("F0_Mean", "acoustic.f0.f0_mean_hz", 2),
    ("F0_SD", "acoustic.f0.f0_sd_hz", 2),
    ("F0_Range", "acoustic.f0.f0_range_hz", 2),
    ("F0_CV", "acoustic.f0.f0_cv", 4),
    ("F0_Slope", "acoustic.f0.f0_slope_hz_per_s", 4),
    # acoustic: voice quality
    ("Jitter_pct", "acoustic.voice_quality.jitter_local_pct", 4),
    ("Shimmer_pct", "acoustic.voice_quality.shimmer_local_pct", 4),
    ("HNR_dB", "acoustic.voice_quality.hnr_mean_db", 2),
    ("H1H2_dB", "acoustic.spectral_tilt.h1_h2_mean_db", 2),
    # acoustic: spectral shape
    ("MFCC1", "acoustic.mfcc.mfcc_1_mean", 2),
    ("MFCC2", "acoustic.mfcc.mfcc_2_mean", 2),
    ("MFCC5", "acoustic.mfcc.mfcc_5_mean", 2),
]

FEATURE_COLUMNS = [c for c, _, _ in FLATTEN]


def _dig(d: dict, dotted: str):
    cur = d
    for part in dotted.split("."):
        if not isinstance(cur, dict):
            return None
        cur = cur.get(part)
    return cur


def flatten_markers(markers: dict, label: str) -> dict:
    """Turn the nested extract_markers() output into one flat row."""
    row = {"Speaker": label}
    for col, path, nd in FLATTEN:
        v = _dig(markers, path)
        if isinstance(v, (int, float)):
            v = round(v, nd) if nd else int(round(v))
        row[col] = v
    return row


def load_chunks(transcript_path: str) -> list[dict]:
    """Read an uploaded transcript JSON into the chunk format extract_markers wants.

    Accepts either ``{"chunks": [...]}`` or a bare list. Each chunk needs a
    ``text`` and a ``timestamp`` with ``start`` / ``end`` (seconds). The ``[*]``
    disfluency markers must be kept: filled-pause and pause markers depend on them.
    """
    with open(transcript_path, encoding="utf-8") as f:
        data = json.load(f)
    raw = data.get("chunks", data) if isinstance(data, dict) else data
    if not isinstance(raw, list):
        raise ValueError("Transcript JSON must be a list of chunks or have a 'chunks' key.")
    chunks = []
    for c in raw:
        ts = c.get("timestamp") or {}
        chunks.append({
            "text": c["text"],
            "timestamp": {"start": ts.get("start"), "end": ts.get("end")},
            "confidence": c.get("confidence"),
        })
    return chunks


def derive_features_with_text(audio_path: str, transcript_path: str, label: str | None = None):
    """Derive features from (audio, transcript); return (flat row, transcript text)."""
    if not label:
        label = os.path.splitext(os.path.basename(audio_path))[0]
    chunks = load_chunks(transcript_path)
    duration_s = round(librosa.get_duration(path=audio_path), 3)
    transcript = {"chunks": chunks, "duration_s": duration_s, "audio_path": audio_path}
    markers = extract_markers(transcript)
    full_text = " ".join(
        c["text"].strip() for c in chunks if c["text"].strip() and c["text"] != "[*]"
    )
    return flatten_markers(markers, label), full_text


def derive_features(audio_path: str, transcript_path: str, label: str | None = None) -> dict:
    """Derive features from (audio, transcript) and return the flat feature row."""
    row, _text = derive_features_with_text(audio_path, transcript_path, label)
    return row