app.py

"""Audio Explorers 2026 — Full Pipeline App.

Two modes:
  - Demo Mode: full pipeline walkthrough (DoA → beamforming → transcribe → gender)
  - Experiment Mode: run preprocessing A/B tests on raw mixtures

Run with:  streamlit run app.py
"""

from __future__ import annotations

import io
import os
import pickle
import tempfile
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt
import librosa
import librosa.display
import soundfile as sf

from src.audio_io import load_mixture
from src.preprocessing import preprocess_enhanced, preprocess_for_ml, preprocess_for_spatial
from src.doa_estimation import estimate_doa, estimate_speaker_directions
from src.post_processing import apply_post_processing_chain
from src.target_enhancement import enhance_target, enhance_all_directions, bilateral_enhance, bilateral_enhance_all, render_binaural, enhance_target_speaker_only, determine_talker_of_interest, _HAS_WPE
from src.transcription import transcribe_detailed
from src.gender_classification import classify_gender_detailed
from src.speaker_diarization import detect_dialogue_pattern, split_transcript_by_speaker, diarize_enhanced, estimate_own_voice_segments, extract_target_speaker
from src.experiment import (
    create_audio_versions,
    run_experiment_on_versions,
    save_demo_results,
)

# params we should use to generate the wave file:
#         presence_gain_db = 3.5
#         comp_threshold_db = -24.0
#         comp_ratio = 3.5
#         target_lufs = -18.0
#         spectral_alpha = 1.45
#         spectral_beta = 0.08
#
#         # De-esser
#         enable_deesser = True
#         deess_threshold_db = -32.0
#         deess_ratio = 5.0
#         deess_freq_low = 4000.0
#         deess_freq_high = 9000.0
#         deess_attack_ms = 2.0
#         deess_release_ms = 60.0

APP_STATE_DIR = Path(".app_state")
DEMO_STATE_PATH = APP_STATE_DIR / "demo_state.pkl"
DEMO_UPLOAD_DIR = APP_STATE_DIR / "demo_uploads"

DEMO_PERSISTED_KEYS = (
    "audio_file",
    "audio_file_label",
    "whisper_model",
    "separation_method",
    "suppression",
    "target_boost",
    "post_processing_enabled",
    "enable_presence",
    "enable_compression",
    "enable_lufs",
    "enable_limiter",
    "enable_spectral_subtract",
    "enable_comfort_noise",
    "use_auto_doa",
    "manual_doas",
    "use_wpe",
    "direction_results",
    "doas",
    "n_speakers",
    "doa_confidences",
    "doa_dominant",
    "doa_overall_confidence",
    "doa_pairwise",
    "doa_meta",
    "doa_raw_candidates",
    "mono_raw",
    "sr_mono",
)


# ---------------------------------------------------------------------------
# Suppression / boost dial helpers
# ---------------------------------------------------------------------------

def _lerp(t: float, low: float, high: float) -> float:
    """Linear interpolation: t=0 → low, t=1 → high."""
    return low + t * (high - low)


def _resolve_suppression(s: float) -> dict:
    """Map the 0–1 Background Suppression slider to DSP parameters."""
    return {
        "mask_sharpness": _lerp(s, 1.5, 5.0),
        "wiener_floor": _lerp(s, 0.15, 0.01),
        "spectral_alpha": _lerp(s, 0.5, 2.5),
        "spectral_beta": _lerp(s, 0.05, 0.005),
        "enable_spectral_subtract": s > 0.3,
    }


def _resolve_target_boost(b: float) -> dict:
    """Map the 0–1 Target Boost slider to post-processing parameters."""
    return {
        "presence_gain_db": _lerp(b, 0.0, 6.0),
        "comp_threshold_db": _lerp(b, -15.0, -30.0),
        "comp_ratio": _lerp(b, 2.0, 8.0),
        "target_lufs": _lerp(b, -23.0, -12.0),
    }


def _ensure_app_state_dirs() -> None:
    APP_STATE_DIR.mkdir(parents=True, exist_ok=True)
    DEMO_UPLOAD_DIR.mkdir(parents=True, exist_ok=True)


def _load_demo_state_from_disk() -> None:
    """Hydrate session state from the last saved demo run once per process."""
    if st.session_state.get("_demo_disk_state_loaded"):
        return

    st.session_state["_demo_disk_state_loaded"] = True
    if not DEMO_STATE_PATH.exists():
        return

    try:
        with DEMO_STATE_PATH.open("rb") as f:
            payload = pickle.load(f)
    except Exception:
        return

    if not isinstance(payload, dict):
        return

    saved_audio = payload.get("audio_file")
    if saved_audio and not Path(saved_audio).exists():
        payload.pop("audio_file", None)
        payload.pop("audio_file_label", None)
        for key in (
            "direction_results",
            "doas",
            "n_speakers",
            "doa_confidences",
            "doa_dominant",
            "doa_overall_confidence",
            "doa_pairwise",
            "doa_meta",
            "doa_raw_candidates",
            "mono_raw",
            "sr_mono",
        ):
            payload.pop(key, None)

    for key in DEMO_PERSISTED_KEYS:
        if key in payload and key not in st.session_state:
            st.session_state[key] = payload[key]


def _save_demo_state_to_disk() -> None:
    """Persist the current demo state so reloads restore the last run."""
    _ensure_app_state_dirs()
    payload = {
        key: st.session_state[key]
        for key in DEMO_PERSISTED_KEYS
        if key in st.session_state
    }
    with DEMO_STATE_PATH.open("wb") as f:
        pickle.dump(payload, f)


def _clear_demo_state() -> None:
    """Remove both in-memory and on-disk demo state."""
    for key in DEMO_PERSISTED_KEYS:
        st.session_state.pop(key, None)
    if DEMO_STATE_PATH.exists():
        DEMO_STATE_PATH.unlink()


def _persist_uploaded_demo_file(uploaded_file) -> tuple[str, str]:
    """Store uploaded demo audio in a stable local path."""
    _ensure_app_state_dirs()
    target = DEMO_UPLOAD_DIR / uploaded_file.name
    target.write_bytes(uploaded_file.getvalue())
    return str(target), uploaded_file.name


# ---------------------------------------------------------------------------
# Page config
# ---------------------------------------------------------------------------

st.set_page_config(
    page_title="Audio Explorers 2026",
    page_icon="🎧",
    layout="wide",
)

# ---------------------------------------------------------------------------
# Tooltip / info content for each pipeline step
# ---------------------------------------------------------------------------

STEP_INFO = {
    "load": (
        "**Loading the raw recording**\n\n"
        "The input is a 4-channel WAV captured by a behind-the-ear (BTE) hearing aid. "
        "The four microphones are: Left-Front (LF), Left-Rear (LR), Right-Front (RF), "
        "Right-Rear (RR). The front-rear spacing on each ear is ~12 mm and the inter-ear "
        "distance is ~17 cm. The recording is at 44.1 kHz."
    ),
    "preprocess": (
        "**Preprocessing**\n\n"
        "Two preprocessing paths are applied:\n"
        "- *Spatial preprocessing* (`preprocess_for_spatial`): highpass filter at 80 Hz to "
        "remove low-frequency rumble, keeps all 4 channels at the original sample rate for "
        "DoA estimation and beamforming.\n"
        "- *ML preprocessing* (`preprocess_for_ml`): highpass → mono downmix → resample to "
        "16 kHz → RMS normalize. Used to produce an audible mono preview of the raw mixture."
    ),
    "doa": (
        "**Direction of Arrival (DoA) estimation**\n\n"
        "Uses a hybrid GCC-PHAT fusion algorithm across 4 microphone pairs "
        "(left same-ear, right same-ear, front cross-ear, rear cross-ear). "
        "The signal is split into short frames; energy-based gating selects "
        "speech-active windows. For each active window and each mic pair, "
        "GCC-PHAT computes the time-delay-of-arrival. These pairwise lag "
        "measurements are fused into an angular score curve (SRP-PHAT–like), "
        "and candidate DoAs are extracted as peaks. Each candidate gets a "
        "confidence score combining peak sharpness, cross-pair agreement, "
        "and active-window coverage."
    ),
    "beamforming": (
        "**Source Separation / Beamforming**\n\n"
        "Three methods are available to isolate each speaker:\n\n"
        "- **MVDR (Minimum Variance Distortionless Response):** Multi-source aware MVDR "
        "beamformer with adaptive diagonal loading, per-interferer covariance estimation, "
        "and Wiener post-filtering. Uses a two-stage approach: separable directions (left/right) "
        "are isolated first, then subtracted so that ambiguous directions (front/back) can be "
        "recovered from the residual. This is the recommended method.\n"
        "- **DAS (Delay-and-Sum):** Aligns channels by compensating propagation "
        "delays using sinc interpolation, then averages. Simple and robust baseline.\n\n"
        "Note: with only ~12 mm front-rear spacing the spatial resolution is limited "
        "for front/back disambiguation, but the two-stage MVDR approach mitigates this."
    ),
    "transcription": (
        "**Automatic Speech Recognition (Whisper)**\n\n"
        "Each separated signal is fed to OpenAI's Whisper model at 16 kHz. "
        "Whisper auto-detects the language and produces a timestamped transcript.\n\n"
        "**Model size matters:** The `base` model often produces gibberish on Danish speech. "
        "Use `medium` or `large` for reliable Danish transcription (at the cost of speed)."
    ),
    "gender": (
        "**Gender Classification (hybrid wav2vec2 + F0)**\n\n"
        "Uses a two-stage approach:\n"
        "1. **wav2vec2 model** (`audeering/wav2vec2-large-robust-6-ft-age-gender`): "
        "a pretrained age/gender classifier. If confidence > 0.85 it decides immediately.\n"
        "2. **F0 (fundamental frequency) analysis** via `librosa.pyin`: median F0 < 165 Hz → male, "
        "≥ 165 Hz → female. Used as fallback or tie-breaker when the model is uncertain.\n\n"
        "The final decision fuses both signals. Returns 'unknown' for very short or unvoiced clips."
    ),
    "results": (
        "**Summary & Speaker Map**\n\n"
        "A polar plot shows detected speakers around the listener. The summary table "
        "lists each speaker's direction, detected language, gender, F0, and a transcript "
        "preview. Expand individual speakers below for full details, audio playback, "
        "and spectrograms."
    ),
}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_spectrogram_fig(audio: np.ndarray, sr: int, title: str) -> plt.Figure:
    if sr != 16000:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
        sr = 16000
    S = librosa.feature.melspectrogram(y=audio.astype(np.float32), sr=sr, n_mels=128)
    S_db = librosa.power_to_db(S, ref=np.max)
    fig, ax = plt.subplots(figsize=(10, 3))
    librosa.display.specshow(S_db, sr=sr, x_axis="time", y_axis="mel", ax=ax)
    ax.set_title(title, fontsize=11)
    plt.tight_layout()
    return fig


def _make_waveform_comparison_fig(
    before: np.ndarray, sr_before: int,
    after: np.ndarray, sr_after: int,
    title: str = "Waveform Comparison",
) -> plt.Figure:
    """Two stacked subplots with RMS envelope for clear before/after comparison."""
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 5), sharex=True,
                                    gridspec_kw={"hspace": 0.15})
    fig.patch.set_facecolor("#fafafa")

    RAW_COLOR = "#e67e22"       # warm orange
    ENHANCED_COLOR = "#27ae60"  # cool green
    ENVELOPE_ALPHA = 0.85

    # Compute RMS envelope
    def _rms_envelope(signal: np.ndarray, sr: int, window_ms: float = 30.0) -> tuple[np.ndarray, np.ndarray]:
        win = max(1, int(sr * window_ms / 1000.0))
        hop = max(1, win // 2)
        n_frames = max(1, (len(signal) - win) // hop + 1)
        rms = np.array([
            np.sqrt(np.mean(signal[i * hop:i * hop + win] ** 2))
            for i in range(n_frames)
        ])
        t = np.array([i * hop / sr for i in range(n_frames)])
        return t, rms

    # Raw mixture (top)
    t_before = np.arange(len(before)) / sr_before
    ax1.fill_between(t_before, before, alpha=0.25, color=RAW_COLOR, linewidth=0)
    ax1.plot(t_before, before, alpha=0.4, linewidth=0.3, color=RAW_COLOR)
    t_env, rms_env = _rms_envelope(before, sr_before)
    ax1.plot(t_env, rms_env, color=RAW_COLOR, linewidth=2, alpha=ENVELOPE_ALPHA, label="RMS envelope")
    ax1.plot(t_env, -rms_env, color=RAW_COLOR, linewidth=2, alpha=ENVELOPE_ALPHA)
    ax1.set_ylabel("Amplitude", fontsize=9)
    ax1.set_title("Raw Mixture (mono)", fontsize=11, fontweight="bold",
                  color=RAW_COLOR, loc="left")
    ax1.grid(True, alpha=0.3, linewidth=0.5)
    ax1.set_facecolor("#fefefe")
    peak_raw = max(np.max(np.abs(before)), 1e-8)
    ax1.set_ylim(-peak_raw * 1.15, peak_raw * 1.15)

    # Enhanced (bottom)
    t_after = np.arange(len(after)) / sr_after
    ax2.fill_between(t_after, after, alpha=0.25, color=ENHANCED_COLOR, linewidth=0)
    ax2.plot(t_after, after, alpha=0.4, linewidth=0.3, color=ENHANCED_COLOR)
    t_env2, rms_env2 = _rms_envelope(after, sr_after)
    ax2.plot(t_env2, rms_env2, color=ENHANCED_COLOR, linewidth=2, alpha=ENVELOPE_ALPHA, label="RMS envelope")
    ax2.plot(t_env2, -rms_env2, color=ENHANCED_COLOR, linewidth=2, alpha=ENVELOPE_ALPHA)
    ax2.set_ylabel("Amplitude", fontsize=9)
    ax2.set_xlabel("Time (s)", fontsize=9)
    ax2.set_title("Enhanced (beamformed)", fontsize=11, fontweight="bold",
                  color=ENHANCED_COLOR, loc="left")
    ax2.grid(True, alpha=0.3, linewidth=0.5)
    ax2.set_facecolor("#fefefe")
    peak_enh = max(np.max(np.abs(after)), 1e-8)
    ax2.set_ylim(-peak_enh * 1.15, peak_enh * 1.15)

    fig.suptitle(title, fontsize=13, fontweight="bold", y=1.01)
    fig.tight_layout()
    return fig


def _generate_pipeline_summary(direction_results, toi, session_state):
    from datetime import datetime, timezone
    lines = ["# Audio Explorers 2026 — Pipeline Summary\n"]
    lines.append(f"*Generated: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M UTC')}*\n")

    # --- Pipeline Configuration ---
    lines.append("## Pipeline Configuration\n")
    lines.append(f"- **Audio file:** {session_state.get('audio_file', 'unknown')}")
    lines.append(f"- **Whisper model:** {session_state.get('whisper_model', 'unknown')}")
    lines.append(f"- **Separation method:** {session_state.get('separation_method', 'unknown')}")
    lines.append(f"- **Suppression level:** {session_state.get('suppression', 'n/a')}")
    lines.append(f"- **Target boost:** {session_state.get('target_boost', 'n/a')}")
    pp_enabled = session_state.get('post_processing_enabled', False)
    lines.append(f"- **Post-processing:** {'enabled' if pp_enabled else 'disabled'}")
    wpe = session_state.get('use_wpe', False)
    lines.append(f"- **WPE dereverberation:** {'enabled' if wpe else 'disabled'}")
    lines.append("")

    # --- DoA Estimation ---
    lines.append("## Direction of Arrival Estimation\n")
    n_speakers = session_state.get('n_speakers', len(direction_results))
    lines.append(f"- **Speakers detected:** {n_speakers}")
    doa_dominant = session_state.get('doa_dominant')
    doa_confidence = session_state.get('doa_overall_confidence')
    if doa_dominant is not None:
        lines.append(f"- **Dominant DoA:** {doa_dominant:.1f}°")
    if doa_confidence is not None:
        lines.append(f"- **Overall confidence:** {doa_confidence:.3f}")

    # ILD info from metadata
    doa_meta = session_state.get('doa_meta', {})
    ild_info = doa_meta.get('ild') if doa_meta else None
    if ild_info and ild_info.get('enabled'):
        lines.append(f"- **ILD (front/back):** "
                     f"{'front' if ild_info.get('global_is_front') else 'rear'} dominant "
                     f"(confidence: {ild_info.get('global_confidence', 0):.3f}, "
                     f"front vote: {ild_info.get('front_vote_fraction', 0):.0%})")

    # Pairwise lag summary
    doa_pairwise = session_state.get('doa_pairwise')
    if doa_pairwise:
        lines.append("\n### Pairwise GCC-PHAT Summary\n")
        lines.append("| Pair | Spacing (mm) | Mean Lag (samples) | σ | Valid % |")
        lines.append("|------|-------------|-------------------|---|---------|")
        for pair_name, info in doa_pairwise.items():
            lines.append(
                f"| {pair_name} | {info['spacing_m']*1000:.1f} | "
                f"{info['lag_mean_samples']:.2f} | {info['lag_std_samples']:.2f} | "
                f"{info['valid_fraction']:.0%} |"
            )

    # Speaker direction table
    doa_confidences = session_state.get('doa_confidences')
    lines.append("\n### Validated Speaker Directions\n")
    lines.append("| # | DoA (°) | Score |")
    lines.append("|---|---------|-------|")
    doas = session_state.get('doas', [r['doa'] for r in direction_results])
    for i, d in enumerate(doas):
        score = f"{doa_confidences[i]:.3f}" if doa_confidences and i < len(doa_confidences) else "n/a"
        lines.append(f"| {i+1} | {d:.1f} | {score} |")
    lines.append("")

    # --- Per-Speaker Analysis ---
    lines.append("## Per-Speaker Analysis\n")
    for i, r in enumerate(direction_results):
        lines.append(f"### Speaker {i+1} — {r['doa']:.0f}°\n")
        f0 = f"{r['median_f0']:.0f} Hz" if not np.isnan(r["median_f0"]) else "n/a"
        lines.append(f"- **Direction:** {r['doa']:.0f}°")
        lines.append(f"- **Language:** {r['language']}")
        lines.append(f"- **Gender:** {r['gender']} (confidence: {r['gender_confidence']:.2f}, method: {r['gender_method']})")
        lines.append(f"- **Median F0:** {f0}")
        dialogue = r.get("dialogue", {})
        if dialogue.get("is_dialogue"):
            lines.append(f"- **⚠️ Dialogue detected:** {dialogue['question_count']} questions — likely 2 speakers (HA user + target)")
        lines.append(f"\n**Full Transcript:**\n")
        lines.append(f"> {r['transcript']}\n")
        ovd_segments = r.get("ovd_segments")
        if ovd_segments and dialogue.get("is_dialogue"):
            lines.append("**Segments (with own-voice detection):**\n")
            lines.append("| Start | End | Speaker | Score | Text |")
            lines.append("|-------|-----|---------|-------|------|")
            for seg in ovd_segments:
                spk = "🎤 HA user" if seg.get("is_own_voice") else "🔊 Target"
                sc = f"{seg.get('own_voice_score', 0):.2f}"
                text = seg.get("text", "")
                if text:
                    lines.append(f"| {seg['start']:.1f}s | {seg['end']:.1f}s | {spk} | {sc} | {text} |")
        elif r.get("segments"):
            lines.append("**Segments:**\n")
            lines.append("| Start | End | Text |")
            lines.append("|-------|-----|------|")
            for seg in r["segments"]:
                lines.append(f"| {seg['start']:.1f}s | {seg['end']:.1f}s | {seg['text']} |")
        lines.append("")

    # --- Talker of Interest ---
    lines.append("## Talker of Interest\n")
    lines.append(f"{toi['reason']}\n")
    if toi.get("score_breakdown"):
        lines.append("### Score Breakdown\n")
        lines.append("| DoA | Total | Frontal | Language | Content | Quality |")
        lines.append("|-----|-------|---------|----------|---------|---------|")
        for s in toi["score_breakdown"]:
            lines.append(
                f"| {s['doa']:.0f}° | {s['total']:.3f} | {s['frontal']:.2f} | "
                f"{s['language']:.2f} | {s['content']:.2f} | {s['quality']:.2f} |"
            )
    lines.append("")

    return "\n".join(lines)


def _audio_bytes(audio: np.ndarray, sr: int) -> bytes:
    buf = io.BytesIO()
    sf.write(buf, audio.astype(np.float32), sr, format="WAV")
    return buf.getvalue()


def _make_polar_fig(
    angles: list[float],
    labels: list[str],
    title: str = "Speaker Directions",
    confidences: list[float] | None = None,
    genders: list[str] | None = None,
    languages: list[str] | None = None,
) -> plt.Figure:
    """Render an informative polar speaker map with head silhouette and badges."""
    SPEAKER_COLORS = ["#e74c3c", "#2ecc71", "#3498db", "#f39c12",
                      "#9b59b6", "#1abc9c", "#e67e22", "#34495e"]
    GENDER_ICONS = {"male": "♂", "female": "♀", "unknown": "?"}

    fig, ax = plt.subplots(subplot_kw={"projection": "polar"}, figsize=(7, 7))
    fig.patch.set_facecolor("#fafafa")
    ax.set_facecolor("#f0f0f0")

    # Head silhouette — circle at center + nose pointing up (front)
    head_angles = np.linspace(0, 2 * np.pi, 80)
    ax.plot(head_angles, np.full_like(head_angles, 0.22), color="#bbb",
            linewidth=2, zorder=2)
    # Nose triangle pointing to 0° (front = North)
    nose_theta = np.array([-0.12, 0.0, 0.12])
    nose_r = np.array([0.22, 0.32, 0.22])
    ax.plot(nose_theta, nose_r, color="#bbb", linewidth=2, zorder=2)
    # Ears
    for ear_sign in [-1, 1]:
        ear_theta = ear_sign * np.pi / 2
        ear_angles = np.linspace(ear_theta - 0.15, ear_theta + 0.15, 20)
        ear_r = 0.22 + 0.04 * np.cos(np.linspace(-np.pi, np.pi, 20))
        ax.plot(ear_angles, ear_r, color="#bbb", linewidth=2, zorder=2)

    # Cardinal labels
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_thetagrids(
        [0, 90, 180, 270],
        labels=["Front\n0°", "Right\n90°", "Back\n180°", "Left\n270°"],
        fontsize=9, fontweight="bold", color="#555",
    )

    # Confidence rings
    ring_levels = [0.25, 0.5, 0.75, 1.0]
    for r in ring_levels:
        ax.plot(np.linspace(0, 2 * np.pi, 100), np.full(100, r),
                color="#ddd", linewidth=0.5, linestyle="--", zorder=1)

    ax.set_ylim(0, 1.55)
    ax.set_yticks([])

    # Plot each speaker
    n = max(len(angles), 1)
    for i, (deg, lbl) in enumerate(zip(angles, labels)):
        theta = np.radians(90 - deg)
        color = SPEAKER_COLORS[i % len(SPEAKER_COLORS)]
        conf = confidences[i] if confidences and i < len(confidences) else 0.75
        radius = 0.4 + 0.6 * conf  # confidence maps to 0.4–1.0

        # Main marker — large filled circle
        marker_size = 16 + 6 * conf
        ax.plot(theta, radius, "o", markersize=marker_size, color=color,
                markeredgecolor="white", markeredgewidth=2, zorder=5)

        # Speaker number inside marker
        ax.text(theta, radius, str(i + 1), ha="center", va="center",
                fontsize=10, fontweight="bold", color="white", zorder=6)

        # Badge: gender icon + language
        gender_str = GENDER_ICONS.get(
            genders[i] if genders and i < len(genders) else "unknown", "?"
        )
        lang_str = languages[i] if languages and i < len(languages) else ""
        badge = f"{gender_str} {lang_str}".strip()

        ax.annotate(
            badge, xy=(theta, radius),
            xytext=(theta, min(radius + 0.22, 1.45)),
            ha="center", va="bottom", fontsize=9, fontweight="bold",
            color=color,
            bbox=dict(boxstyle="round,pad=0.2", facecolor="white",
                      edgecolor=color, alpha=0.85),
            zorder=7,
        )

    # Legend
    from matplotlib.lines import Line2D
    legend_handles = []
    for i, deg in enumerate(angles):
        color = SPEAKER_COLORS[i % len(SPEAKER_COLORS)]
        gender_str = (genders[i] if genders and i < len(genders) else "?")
        lang_str = (languages[i] if languages and i < len(languages) else "")
        conf_str = f"{confidences[i]:.0%}" if confidences and i < len(confidences) else ""
        lbl = f"Spk {i+1}: {deg:.0f}° · {gender_str} · {lang_str} ({conf_str})"
        legend_handles.append(
            Line2D([0], [0], marker="o", color="w", markerfacecolor=color,
                   markersize=10, label=lbl)
        )
    if legend_handles:
        ax.legend(handles=legend_handles, loc="upper left",
                  bbox_to_anchor=(-0.15, -0.05), fontsize=8,
                  frameon=True, fancybox=True, shadow=True)

    ax.set_title(title, pad=30, fontsize=14, fontweight="bold")
    fig.tight_layout()
    return fig


def _step_note(key: str) -> None:
    """Render an expandable info note for a pipeline step."""
    with st.expander("ℹ️ How does this work?", expanded=False):
        st.markdown(STEP_INFO[key])


# ---------------------------------------------------------------------------
# Sidebar — mode selection
# ---------------------------------------------------------------------------

mode = st.sidebar.radio(
    "Mode",
    ["Demo", "Enhancement Only", "Experiment", "Post-processing Playground"],  # ✅ ADDED MODE
    index=0,
    help="Demo: full pipeline walkthrough. Enhancement Only: run enhancement only. Experiment: preprocessing A/B tests.",
)


# ===================================================================
# DEMO MODE — Full pipeline
# ===================================================================

if mode == "Demo":
    _load_demo_state_from_disk()
    st.title("Audio Explorers 2026")
    st.markdown(
        "Extracting the talker of interest from a multi-talker hearing-aid recording. "
        "Walk through the full pipeline step by step: "
        "**load → preprocess → DoA estimation → beamforming → transcription → gender classification.**"
    )

    # --- Sidebar settings ---
    with st.sidebar:
        st.header("Settings")
        uploaded_file = st.file_uploader(
            "Upload a 4-channel WAV file",
            type=["wav"],
            help="Upload your own multi-channel hearing-aid recording. Must be a 4-channel WAV file.",
        )
        default_audio_options = ["data/raw/example_mixture.wav", "data/raw/mixture.wav"]
        saved_audio_file = st.session_state.get("audio_file")
        saved_audio_label = st.session_state.get("audio_file_label")
        if uploaded_file is not None:
            audio_file, audio_label = _persist_uploaded_demo_file(uploaded_file)
            st.session_state["audio_file"] = audio_file
            st.session_state["audio_file_label"] = audio_label
            st.sidebar.success(f"Uploaded: {audio_label}")
        else:
            default_index = 0
            if saved_audio_file in default_audio_options:
                default_index = default_audio_options.index(saved_audio_file)
            audio_file = st.selectbox(
                "Audio file",
                default_audio_options,
                index=default_index,
            )
            audio_label = Path(audio_file).name
            if saved_audio_file and saved_audio_file not in default_audio_options and Path(saved_audio_file).exists():
                audio_file = saved_audio_file
                audio_label = saved_audio_label or Path(saved_audio_file).name
                st.caption(f"Reloaded saved uploaded file: `{audio_label}`")
            st.session_state["audio_file"] = audio_file
            st.session_state["audio_file_label"] = audio_label
        whisper_model = st.selectbox(
            "Whisper model",
            ["tiny", "base", "small", "medium", "large"],
            index=["tiny", "base", "small", "medium", "large"].index(
                st.session_state.get("whisper_model", "medium")
            ),
            help="'medium' or 'large' recommended for Danish speech.",
        )
        separation_method = st.selectbox(
            "Separation method",
            ["mvdr", "das"],
            index=["mvdr", "das"].index(st.session_state.get("separation_method", "mvdr")),
            help=(
                "MVDR: multi-source aware MVDR with adaptive loading and Wiener post-filter (recommended). "
                "DAS: delay-and-sum beamforming (fast baseline)."
            ),
        )
        st.divider()
        suppression = st.slider(
            "🔇 Background Suppression",
            min_value=0.0, max_value=1.0, value=float(st.session_state.get("suppression", 0.5)), step=0.05,
            help=(
                "How aggressively to suppress background speakers and noise. "
                "Controls spatial mask sharpness, Wiener post-filter depth, and spectral subtraction. "
                "Higher = quieter background but may introduce artifacts."
            ),
        )
        target_boost = st.slider(
            "🔊 Target Boost",
            min_value=0.0, max_value=1.0, value=float(st.session_state.get("target_boost", 0.5)), step=0.05,
            help=(
                "How much to boost the target speaker's clarity and loudness. "
                "Controls presence boost (2–6 kHz), dynamic compression, and LUFS normalization. "
                "Higher = louder and clearer target."
            ),
        )
        sup_params = _resolve_suppression(suppression)
        boost_params = _resolve_target_boost(target_boost)
        st.divider()
        enable_post_processing = st.toggle(
            "Apply post-processing",
            value=bool(st.session_state.get("post_processing_enabled", True)),
            help=(
                "Apply presence boost, compression, LUFS normalization, and limiting "
                "to the beamformed mono outputs used for listening and download."
            ),
        )
        with st.expander("Advanced post-processing overrides", expanded=False):
            st.caption("These are auto-set by the sliders above. Override only if needed.")
            enable_presence = st.checkbox("Presence boost", value=bool(st.session_state.get("enable_presence", True)))
            enable_compression = st.checkbox("Compression", value=bool(st.session_state.get("enable_compression", True)))
            enable_lufs = st.checkbox("LUFS normalization", value=bool(st.session_state.get("enable_lufs", True)))
            enable_limiter = st.checkbox("Limiter", value=bool(st.session_state.get("enable_limiter", True)))
            enable_spectral_subtract = st.checkbox(
                "Spectral subtraction",
                value=bool(st.session_state.get("enable_spectral_subtract", sup_params["enable_spectral_subtract"])),
                help="Auto-enabled when Background Suppression > 0.3. Can be aggressive — use with care.",
            )
            enable_comfort_noise = st.checkbox("Comfort noise", value=bool(st.session_state.get("enable_comfort_noise", False)),
                help="Add low-level shaped noise to prevent uncomfortable silence after noise reduction.")
        use_auto_doa = st.toggle(
            "Auto-detect DoAs",
            value=bool(st.session_state.get("use_auto_doa", True)),
            help="Use the GCC-PHAT fusion algorithm to find speaker directions automatically.",
        )
        if not use_auto_doa:
            st.caption(
                "0° = front, 90° = right, 180° = back, 270° = left."
            )
            doa_text = st.text_input(
                "Comma-separated DoAs",
                value=", ".join(str(x) for x in st.session_state.get("manual_doas", [0, 90, 180, 270])),
            )
            try:
                manual_doas = [float(x.strip()) for x in doa_text.split(",") if x.strip()]
            except ValueError:
                st.error("Enter numeric angles separated by commas.")
                st.stop()

        # WPE dereverberation toggle
        use_wpe = st.toggle(
            "WPE Dereverberation (experimental)",
            value=bool(st.session_state.get("use_wpe", False)),
            help=(
                "Apply Weighted Prediction Error dereverberation before MVDR beamforming. "
                "Reduces room reverberation but is memory-intensive. Requires nara-wpe."
            ),
        )
        if use_wpe and not _HAS_WPE:
            st.warning("⚠️ nara-wpe is not installed. WPE dereverberation will be skipped.")
        if st.button("Clear saved demo state", use_container_width=True):
            _clear_demo_state()
            st.rerun()

    st.session_state["audio_file"] = audio_file
    st.session_state["audio_file_label"] = audio_label
    st.session_state["whisper_model"] = whisper_model
    st.session_state["separation_method"] = separation_method
    st.session_state["suppression"] = suppression
    st.session_state["target_boost"] = target_boost
    st.session_state["post_processing_enabled"] = enable_post_processing
    st.session_state["enable_presence"] = enable_presence
    st.session_state["enable_compression"] = enable_compression
    st.session_state["enable_lufs"] = enable_lufs
    st.session_state["enable_limiter"] = enable_limiter
    st.session_state["enable_spectral_subtract"] = enable_spectral_subtract
    st.session_state["enable_comfort_noise"] = enable_comfort_noise
    st.session_state["use_auto_doa"] = use_auto_doa
    st.session_state["manual_doas"] = manual_doas if not use_auto_doa else []
    st.session_state["use_wpe"] = use_wpe
    _save_demo_state_to_disk()

    # --- Load ---
    if not Path(audio_file).exists():
        st.error(f"File not found: `{audio_file}`")
        st.stop()

    audio_raw, sr = load_mixture(audio_file)

    # Validate 4-channel audio
    if audio_raw.ndim != 2 or audio_raw.shape[1] != 4:
        st.error(
            f"Expected a 4-channel WAV file, got shape {audio_raw.shape}. "
            "The file must have exactly 4 channels in order: "
            "[left front, left rear, right front, right rear]."
        )
        st.stop()

    # ================================================================
    # Step 1 — Raw mixture
    # ================================================================
    st.header("1  Load Raw Recording")
    _step_note("load")
    col1, col2 = st.columns([1, 2])
    with col1:
        st.markdown(
            f"**Channels:** {audio_raw.shape[1]}  \n"
            f"**Sample rate:** {sr} Hz  \n"
            f"**Duration:** {audio_raw.shape[0]/sr:.1f} s"
        )
        mono_raw, sr_mono = preprocess_for_ml(audio_raw, sr)
        st.audio(_audio_bytes(mono_raw, sr_mono), format="audio/wav")
    with col2:
        fig = _make_spectrogram_fig(mono_raw, sr_mono, "Raw mixture (mono downmix, 16 kHz)")
        st.pyplot(fig)
        plt.close(fig)

    # ================================================================
    # Step 2 — Preprocessing
    # ================================================================
    st.header("2  Preprocessing")
    _step_note("preprocess")
    st.caption(
        "Spatial preprocessing keeps all 4 channels at 44.1 kHz (for DoA + beamforming). "
        "ML preprocessing produces a mono 16 kHz signal (for the player above)."
    )

    # ================================================================
    # Run pipeline
    # ================================================================
    if st.button("Run full pipeline", type="primary"):

        # -- Spatial preprocess --
        audio_spatial, sr_spatial = preprocess_for_spatial(audio_raw, sr)

        # -- DoA estimation --
        if use_auto_doa:
            with st.spinner("Estimating directions of arrival..."):
                doa_result = estimate_doa(audio_spatial, sr_spatial)
                # Use validated speaker directions (filtered + clustered)
                dir_result = estimate_speaker_directions(audio_raw, sr)
            n_speakers = dir_result["count"]
            directions = dir_result["directions"]
            doas = [d["doa"] for d in directions]
            doa_confidences = [d["score"] for d in directions]
            doa_dominant = doas[0] if doas else 0.0
            doa_overall_confidence = doa_result["confidence"]
            doa_pairwise = doa_result["pairwise_lags"]
            doa_meta = doa_result["metadata"]
            doa_raw_candidates = doa_result["candidate_doas"]
        else:
            doas = manual_doas
            doa_confidences = None
            doa_dominant = doas[0] if doas else 0.0
            doa_overall_confidence = None
            doa_pairwise = None
            doa_meta = None
            doa_raw_candidates = None
            n_speakers = len(doas)

        if not doas:
            st.warning("No speaker directions found. Try lowering the energy threshold or providing manual DoAs.")
            st.stop()

        # --- Hardcoded tuning (final values) ---
        presence_gain_db = 3.5
        comp_threshold_db = -24.0
        comp_ratio = 3.5
        target_lufs = -18.0
        spectral_alpha = 1.45
        spectral_beta = 0.08

        # De-esser
        enable_deesser = True
        deess_threshold_db = -32.0
        deess_ratio = 5.0
        deess_freq_low = 4000.0
        deess_freq_high = 9000.0
        deess_attack_ms = 2.0
        deess_release_ms = 60.0

        # -- Source separation / beamforming --
        beamform_method = separation_method  # "mvdr" or "das"
        with st.spinner(f"Beamforming {len(doas)} directions ({beamform_method.upper()})..."):
            enhanced_raw = enhance_all_directions(
                audio_spatial, sr_spatial, doas, method=beamform_method,
                use_wpe=use_wpe,
                mask_sharpness=sup_params["mask_sharpness"],
                wiener_floor=sup_params["wiener_floor"],
            )

        sr_enhanced_raw = sr_spatial

        # --- Post-processing ---
        if enable_post_processing:
            processed_raw = [
                apply_post_processing_chain(
                    sig,
                    sr_spatial,
                    enable_presence=enable_presence,
                    enable_compression=enable_compression,
                    enable_lufs=enable_lufs,
                    enable_limiter=enable_limiter,
                    enable_spectral_subtract=enable_spectral_subtract,
                    enable_comfort_noise=enable_comfort_noise,

                    # Hardcoded tuning
                    presence_gain_db=presence_gain_db,
                    comp_threshold_db=comp_threshold_db,
                    comp_ratio=comp_ratio,
                    target_lufs=target_lufs,
                    spectral_alpha=spectral_alpha,
                    spectral_beta=spectral_beta,

                    # De-esser
                    enable_deesser=enable_deesser,
                    deess_threshold_db=deess_threshold_db,
                    deess_ratio=deess_ratio,
                )
                for sig in enhanced_raw
            ]
        else:
            processed_raw = [sig.copy() for sig in enhanced_raw]

        stereo_raw = [
            render_binaural(sig, sr_spatial, doa, raw_multichannel=audio_raw)
            for sig, doa in zip(processed_raw, doas)
        ]


        enhanced_ml = []
        for sig in enhanced_raw:
            sig_ml, _ = preprocess_enhanced(sig, sr_spatial)
            enhanced_ml.append(sig_ml)

        # -- Transcription + Gender --
        direction_results = []
        for i, doa in enumerate(doas):
            with st.spinner(f"Analysing speaker at {doa:.0f}°..."):
                asr = transcribe_detailed(
                    enhanced_ml[i].astype(np.float32), 16000,
                    model_size=whisper_model,
                )
                gen = classify_gender_detailed(
                    enhanced_ml[i].astype(np.float32), 16000,
                )
                # Dialogue detection
                dialogue = detect_dialogue_pattern(asr["segments"])

                # Own-voice detection (uses raw 4-channel + enhanced mono)
                ovd_segments = None
                audio_target_only = None
                if dialogue["is_dialogue"]:
                    ovd_segments = estimate_own_voice_segments(
                        asr["segments"], audio_raw, sr,
                        enhanced_mono=enhanced_ml[i],
                        enhanced_sr=16000,
                    )
                    # Generate target-speaker-only audio
                    target_ts = [
                        (s["start"], s["end"])
                        for s in ovd_segments
                        if not s.get("is_own_voice") and s.get("text", "").strip()
                    ]
                    if target_ts:
                        audio_target_only = enhance_target_speaker_only(
                            audio_spatial, sr_spatial,
                            target_doa=doa,
                            target_timestamps=target_ts,
                            doa_list=doas,
                            method=beamform_method,
                            use_wpe=use_wpe,
                        )

                direction_results.append({
                    "doa": doa,
                    "audio_ml": enhanced_ml[i],
                    "sr_ml": 16000,
                    "audio_post": processed_raw[i],
                    "sr_post": sr_enhanced_raw,
                    "audio_beamformed": enhanced_raw[i],
                    "sr_beamformed": sr_enhanced_raw,
                    "audio_stereo": stereo_raw[i],
                    "sr_stereo": sr_enhanced_raw,
                    "audio_target_only": audio_target_only,
                    "sr_target_only": sr_spatial,
                    "transcript": asr["text"],
                    "language": asr["language"],
                    "segments": asr["segments"],
                    "gender": gen["gender"],
                    "gender_confidence": gen["confidence"],
                    "gender_method": gen["method"],
                    "median_f0": gen["median_f0"],
                    "dialogue": dialogue,
                    "ovd_segments": ovd_segments,
                })

        # Update speaker count: DOA directions + intra-DOA dialogue splits
        n_dialogue_extras = sum(
            1 for r in direction_results if r["dialogue"]["is_dialogue"]
        )
        n_speakers = n_speakers + n_dialogue_extras

        # Store in session state
        st.session_state["direction_results"] = direction_results
        st.session_state["doas"] = doas
        st.session_state["n_speakers"] = n_speakers
        st.session_state["doa_confidences"] = doa_confidences
        st.session_state["doa_dominant"] = doa_dominant
        st.session_state["doa_overall_confidence"] = doa_overall_confidence
        st.session_state["doa_pairwise"] = doa_pairwise
        st.session_state["doa_meta"] = doa_meta
        st.session_state["doa_raw_candidates"] = doa_raw_candidates
        st.session_state["mono_raw"] = mono_raw
        st.session_state["sr_mono"] = sr_mono
        st.session_state["separation_method"] = separation_method
        st.session_state["post_processing_enabled"] = enable_post_processing
        st.session_state["whisper_model"] = whisper_model
        st.session_state["use_wpe"] = use_wpe
        st.session_state["suppression"] = suppression
        st.session_state["target_boost"] = target_boost
        st.session_state["audio_file"] = audio_file
        st.session_state["audio_file_label"] = audio_label
        st.session_state["enable_presence"] = enable_presence
        st.session_state["enable_compression"] = enable_compression
        st.session_state["enable_lufs"] = enable_lufs
        st.session_state["enable_limiter"] = enable_limiter
        st.session_state["enable_spectral_subtract"] = enable_spectral_subtract
        st.session_state["enable_comfort_noise"] = enable_comfort_noise
        st.session_state["use_auto_doa"] = use_auto_doa
        st.session_state["manual_doas"] = manual_doas if not use_auto_doa else []
        _save_demo_state_to_disk()


    # ================================================================
    # Display results (persisted in session state)
    # ================================================================
    if "direction_results" not in st.session_state:
        st.info("Press **Run full pipeline** above to start processing.")
        st.stop()

    direction_results = st.session_state["direction_results"]
    doas = st.session_state["doas"]
    n_speakers = st.session_state.get("n_speakers", len(doas))
    doa_confidences = st.session_state.get("doa_confidences")
    doa_dominant = st.session_state.get("doa_dominant")
    doa_overall_confidence = st.session_state.get("doa_overall_confidence")
    doa_pairwise = st.session_state.get("doa_pairwise")
    doa_meta = st.session_state.get("doa_meta")
    doa_raw_candidates = st.session_state.get("doa_raw_candidates")

    # sort the results from 0 to 360, for consistency. (also a sane default of speaker 0deg being first)
    direction_results = sorted(direction_results, key=lambda r: r["doa"])


    # ================================================================
    # Step 3 — DoA estimation results
    # ================================================================
    st.header("3  Direction of Arrival (DoA)")
    _step_note("doa")

    if doa_overall_confidence is not None:
        st.success(
            f"Detected **{n_speakers} speaker(s)** "
            f"({len(doas)} spatial directions"
            + (f", +{n_speakers - len(doas)} from dialogue detection" if n_speakers > len(doas) else "")
            + f"). Dominant DoA: **{doa_dominant:.1f}°** "
            f"(confidence: {doa_overall_confidence:.2f})"
        )
        doas = sorted(doas)

        # Validated speaker direction table
        cand_rows = []
        for i, d in enumerate(doas):
            row = {"#": i + 1, "DoA (°)": f"{d:.1f}"}
            if doa_confidences and i < len(doa_confidences):
                row["Score"] = f"{doa_confidences[i]:.3f}"
            cand_rows.append(row)
        st.dataframe(pd.DataFrame(cand_rows), use_container_width=False, hide_index=True)

        # sort the results from 0 to 360, for consistency. (also a sane default of speaker 0deg being first)
        if doa_raw_candidates:
            doa_raw_candidates = sorted(doa_raw_candidates, key=lambda r: r["doa"])
        # Show raw candidates in collapsed section
        if doa_raw_candidates:
            with st.expander(f"Raw DoA candidates ({len(doa_raw_candidates)} before filtering)"):
                raw_rows = []
                for i, c in enumerate(doa_raw_candidates):
                    raw_rows.append({
                        "#": i + 1,
                        "DoA (°)": f"{c['doa']:.1f}",
                        "Score": f"{c['score']:.3f}",
                        "Support pairs": ", ".join(c.get("support_pairs", [])),
                    })
                st.dataframe(pd.DataFrame(raw_rows), use_container_width=False, hide_index=True)
                st.caption(
                    f"Filtering: min_score ≥ 0.65, min_support_pairs ≥ 2, "
                    f"cluster separation ≥ 35°. "
                    f"Reduced {len(doa_raw_candidates)} raw peaks → {n_speakers} validated speakers."
                )

        # Pairwise lag details (collapsed)
        if doa_pairwise:
            with st.expander("Pairwise GCC-PHAT details"):
                for pair_name, info in doa_pairwise.items():
                    st.markdown(
                        f"**{pair_name}** — channels {info['channels']}, "
                        f"spacing {info['spacing_m']*1000:.1f} mm, "
                        f"active windows: {info['active_window_count']}, "
                        f"valid: {info['valid_window_count']} "
                        f"({info['valid_fraction']:.0%}), "
                        f"mean lag: {info['lag_mean_samples']:.2f} samples "
                        f"(σ={info['lag_std_samples']:.2f})"
                    )
    else:
        st.info(f"Using **manual DoAs**: {', '.join(f'{d:.0f}°' for d in doas)}")

    # ================================================================
    # Step 4 — Source separation
    # ================================================================
    active_method = st.session_state.get("separation_method", separation_method)
    st.header(f"4  Beamforming ({active_method.upper()})")
    _step_note("beamforming")
    st.caption(
        f"Each detected direction is isolated using {active_method.upper()} beamforming."
    )
    if st.session_state.get("post_processing_enabled", False):
        st.caption("Playback and downloads use the post-processing chain after beamforming.")

    # ================================================================
    # Step 5 — Transcription
    # ================================================================
    st.header(f"5  Transcription (Whisper {st.session_state.get('whisper_model', whisper_model)})")
    _step_note("transcription")

    # ================================================================
    # Step 6 — Gender classification
    # ================================================================
    st.header("6  Gender Classification")
    _step_note("gender")

    # ================================================================
    # Step 7 — Summary & speaker map
    # ================================================================
    st.header("7  Results")
    _step_note("results")

    # Polar plot
    polar_labels = []
    polar_genders = []
    polar_languages = []
    for r in direction_results:
        lbl = f"{r['doa']:.0f}°\n{r['gender']}"
        if not np.isnan(r["median_f0"]):
            lbl += f"\nF0={r['median_f0']:.0f}Hz"
        lbl += f"\n{r['language']}"
        polar_labels.append(lbl)
        polar_genders.append(r["gender"])
        polar_languages.append(r["language"])

    col1, col2 = st.columns([1, 2])
    with col1:
        fig_polar = _make_polar_fig(
            doas, polar_labels, "Detected Speakers",
            confidences=[c / max(doa_confidences) if doa_confidences and max(doa_confidences) > 0 else 1.0
                         for c in (doa_confidences or [1.0] * len(doas))],
            genders=polar_genders,
            languages=polar_languages,
        )
        st.pyplot(fig_polar)
        plt.close(fig_polar)
    with col2:
        summary_rows = []
        for r in direction_results:
            dialogue = r.get("dialogue", {})
            spk_label = "🗣️🗣️ 2 speakers" if dialogue.get("is_dialogue") else "🗣️ 1 speaker"
            summary_rows.append({
                "Direction": f"{r['doa']:.0f}°",
                "Speakers": spk_label,
                "Language": r["language"],
                "Gender": r["gender"],
                "Gender Conf.": f"{r['gender_confidence']:.2f}",
                "Median F0": f"{r['median_f0']:.0f} Hz" if not np.isnan(r["median_f0"]) else "n/a",
                "Transcript": r["transcript"][:100] + ("…" if len(r["transcript"]) > 100 else ""),
            })
        st.dataframe(pd.DataFrame(summary_rows), use_container_width=True, hide_index=True)

    # ================================================================
    # Per-speaker details
    # ================================================================
    st.header("8  Per-Speaker Details")
    for r in direction_results:
        with st.expander(
            f"Speaker at {r['doa']:.0f}° — {r['gender']} — {r['language']}",
            expanded=(r["doa"] == doa_dominant),
        ):
            col1, col2 = st.columns([1, 2])
            with col1:
                st.markdown(f"**Direction:** {r['doa']:.0f}°")
                st.markdown(f"**Language:** {r['language']}")
                st.metric("Gender", r["gender"])
                st.progress(min(r["gender_confidence"], 1.0))
                st.caption(
                    f"Confidence: {r['gender_confidence']:.2f} · "
                    f"F0: {r['median_f0']:.0f} Hz · "
                    f"Method: {r['gender_method']}"
                )
                st.audio(_audio_bytes(r["audio_post"], r["sr_post"]), format="audio/wav")
                if "audio_stereo" in r:
                    st.caption("🎧 Stereo (bilateral beamforming)")
                    st.audio(_audio_bytes(r["audio_stereo"], r["sr_stereo"]), format="audio/wav")
                    st.download_button(
                        f"⬇ Download stereo WAV ({r['doa']:.0f}°)",
                        _audio_bytes(r["audio_stereo"], r["sr_stereo"]),
                        f"stereo_{r['doa']:.0f}deg.wav",
                        "audio/wav",
                        key=f"dl_stereo_{r['doa']:.0f}",
                    )
            with col2:
                fig = _make_spectrogram_fig(
                    r["audio_post"], r["sr_post"], f"Enhanced — {r['doa']:.0f}°"
                )
                st.pyplot(fig)
                plt.close(fig)

            st.markdown("**Transcript:**")

            dialogue = r.get("dialogue", {})
            ovd_segments = r.get("ovd_segments")
            if dialogue.get("is_dialogue"):
                st.warning(
                    f"🗣️🗣️ **Dialogue detected** — {dialogue['question_count']} questions found. "
                    f"This direction likely contains **2 speakers** (e.g., HA user + conversation partner)."
                )
                st.caption(dialogue.get("evidence", ""))

                # Show target-speaker-only audio if available
                if "audio_target_only" in r:
                    st.caption("🔊 **Target speaker only** (HA user's voice reduced)")
                    st.audio(_audio_bytes(r["audio_target_only"], r["sr_target_only"]), format="audio/wav")
                    st.download_button(
                        f"⬇ Download target-speaker-only WAV ({r['doa']:.0f}°)",
                        _audio_bytes(r["audio_target_only"], r["sr_target_only"]),
                        f"target_only_{r['doa']:.0f}deg.wav",
                        "audio/wav",
                        key=f"dl_target_only_{r['doa']:.0f}",
                    )

            st.info(r["transcript"])

            if r["segments"]:
                st.markdown("**Segments:**")

                # Use OVD results if available, otherwise basic display
                if ovd_segments and dialogue.get("is_dialogue"):
                    n_ov = sum(1 for s in ovd_segments if s.get("is_own_voice"))
                    n_target = len(ovd_segments) - n_ov
                    st.caption(f"🎤 = likely HA user ({n_ov} segments) · 🔊 = likely target speaker ({n_target} segments)")
                    for seg in ovd_segments:
                        if seg.get("is_own_voice"):
                            prefix = "🎤"
                            score_str = f" (OVD={seg['own_voice_score']:.2f})"
                        else:
                            prefix = "🔊"
                            score_str = ""
                        text = seg.get("text", "")
                        if text:
                            st.caption(f"{prefix} [{seg['start']:.1f}s – {seg['end']:.1f}s]{score_str} {text}")
                else:
                    for seg in r["segments"]:
                        text = seg["text"]
                        st.caption(f"💬 [{seg['start']:.1f}s – {seg['end']:.1f}s] {text}")

    # ================================================================
    # Talker of interest
    # ================================================================
    st.header("9  Talker of Interest")
    toi = determine_talker_of_interest(direction_results)
    st.success(toi["reason"])

    if toi["score_breakdown"]:
        with st.expander("Score breakdown"):
            toi_rows = []
            for s in toi["score_breakdown"]:
                toi_rows.append({
                    "DoA": f"{s['doa']:.0f}°",
                    "Total": f"{s['total']:.3f}",
                    "Frontal": f"{s['frontal']:.2f}",
                    "Language": f"{s['language']:.2f}",
                    "Content": f"{s['content']:.2f}",
                    "Quality": f"{s['quality']:.2f}",
                })
            st.dataframe(pd.DataFrame(toi_rows), use_container_width=False, hide_index=True)

    # ================================================================
    # Before / After comparison
    # ================================================================
    st.header("10  Before vs After")
    st.markdown("Compare the raw mixture with the separated/beamformed direction.")
    pick_dir = st.selectbox(
        "Compare with direction",
        [f"{r['doa']:.0f}°" for r in direction_results],
    )
    picked = next(r for r in direction_results if f"{r['doa']:.0f}°" == pick_dir)

    col1, col2 = st.columns(2)
    with col1:
        st.markdown("**Raw mixture (mono)**")
        st.audio(
            _audio_bytes(st.session_state["mono_raw"], st.session_state["sr_mono"]),
            format="audio/wav",
        )
        fig_before = _make_spectrogram_fig(
            st.session_state["mono_raw"], st.session_state["sr_mono"], "Raw mixture"
        )
        st.pyplot(fig_before)
        plt.close(fig_before)
    with col2:
        method_label = st.session_state.get("separation_method", "enhanced")
        if st.session_state.get("post_processing_enabled", False):
            method_label = f"{method_label} + post"
        st.markdown(f"**Enhanced — {pick_dir} ({method_label})**")
        st.audio(_audio_bytes(picked["audio_post"], picked["sr_post"]), format="audio/wav")
        fig_after = _make_spectrogram_fig(
            picked["audio_post"], picked["sr_post"], f"Enhanced {pick_dir}"
        )
        st.pyplot(fig_after)
        plt.close(fig_after)

    fig_wave = _make_waveform_comparison_fig(
        st.session_state["mono_raw"], st.session_state["sr_mono"],
        picked["audio_post"], picked["sr_post"],
    )
    st.pyplot(fig_wave)
    plt.close(fig_wave)

    # ================================================================
    # Target speaker selection + download
    # ================================================================
    st.header("11  Select Target Speaker")
    st.markdown(
        "Pick the talker of interest and download the enhanced WAV file. "
        "The auto-selected target is based on frontal direction, language, and content analysis."
    )
    target_dir = st.selectbox(
        "Target speaker direction",
        [f"{r['doa']:.0f}°" for r in direction_results],
        key="target_select",
    )
    target = next(r for r in direction_results if f"{r['doa']:.0f}°" == target_dir)

    col1, col2 = st.columns(2)
    with col1:
        st.caption("🔊 Mono enhanced")
        st.audio(_audio_bytes(target["audio_post"], target["sr_post"]), format="audio/wav")
        st.download_button(
            f"Download mono WAV ({target_dir})",
            _audio_bytes(target["audio_post"], target["sr_post"]),
            f"enhanced_{target['doa']:.0f}deg.wav",
            "audio/wav",
            key="dl_target_mono",
        )
        if "audio_stereo" in target:
            st.caption("🎧 Stereo (bilateral beamforming)")
            st.audio(_audio_bytes(target["audio_stereo"], target["sr_stereo"]), format="audio/wav")
            st.download_button(
                f"Download stereo WAV ({target_dir})",
                _audio_bytes(target["audio_stereo"], target["sr_stereo"]),
                f"enhanced_stereo_{target['doa']:.0f}deg.wav",
                "audio/wav",
                key="dl_target_stereo",
            )
    with col2:
        st.markdown(f"**Gender:** {target['gender']} ({target['gender_confidence']:.2f})")
        st.markdown(f"**Language:** {target['language']}")
        st.markdown(f"**Transcript:** {target['transcript']}")

    # ================================================================
    # Pipeline summary export
    # ================================================================
    st.header("12  Export Summary")
    summary_md = _generate_pipeline_summary(direction_results, toi, st.session_state)
    st.download_button("Download pipeline summary", summary_md, "pipeline_summary.md", "text/markdown")

# ================================================================
# 🎛 POST-PROCESSING PLAYGROUND (SIMPLIFIED)
# ================================================================
elif mode == "Post-processing Playground":

    st.title("Post-processing Playground")

    # ------------------------------------------------------------
    # INPUT + CONTROLS
    # ------------------------------------------------------------
    with st.sidebar:
        st.header("Input")

        uploaded_file = st.file_uploader(
            "Upload enhanced audio (mono or stereo WAV)",
            type=["wav"]
        )

        st.divider()
        st.header("Processing")

        enable_presence = st.checkbox("Presence boost", True)
        enable_compression = st.checkbox("Compression", True)
        enable_lufs = st.checkbox("LUFS normalization", True)
        enable_limiter = st.checkbox("Limiter", True)
        enable_spectral_subtract = st.checkbox("Spectral subtraction", False)  # keep OFF (important)
        enable_comfort_noise = st.checkbox("Comfort noise", False)

        st.divider()
        st.header("Tuning")

        # --- Presence (LESS aggressive) ---
        presence_gain_db = st.slider("Presence gain (dB)", 0.0, 12.0, 3.5, 0.5)

        # --- De-esser ---
        deess_threshold_db = st.slider("De-esser threshold (dB)", -50.0, -10.0, -32.0, 1.0)
        deess_ratio = st.slider("De-esser ratio", 1.0, 10.0, 5.0, 0.5)

        # --- Compression (stronger for LF spikes) ---
        comp_threshold_db = st.slider("Compression threshold (dB)", -40.0, -5.0, -24.0, 1.0)
        comp_ratio = st.slider("Compression ratio", 1.0, 8.0, 3.5, 0.1)

        target_lufs = st.slider("Target LUFS", -30.0, -10.0, -18.0, 1.0)

        # --- Spectral subtraction (leave mild) ---
        spectral_alpha = st.slider("Spectral alpha", 0.8, 2.0, 1.1, 0.05)
        spectral_beta = st.slider("Spectral beta", 0.01, 0.5, 0.08, 0.01)


    # ------------------------------------------------------------
    # LOAD AUDIO
    # ------------------------------------------------------------
    if uploaded_file is None:
        st.info("Upload an enhanced audio file to begin.")
        st.stop()

    import soundfile as sf
    import tempfile

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp.write(uploaded_file.getvalue())
        temp_path = tmp.name

    audio, sr = sf.read(temp_path)

    # Convert stereo → mono if needed
    if audio.ndim == 2:
        audio = np.mean(audio, axis=1)

    st.subheader("🔊 Input Audio")
    st.audio(_audio_bytes(audio, sr), format="audio/wav")

    # ------------------------------------------------------------
    # APPLY SETTINGS (ONLY BUTTON YOU NEED)
    # ------------------------------------------------------------
    if st.button("🎧 Apply Post-processing", type="primary"):

        with st.spinner("Processing..."):

            processed = apply_post_processing_chain(
                audio,
                sr,
                enable_presence=enable_presence,
                enable_compression=enable_compression,
                enable_lufs=enable_lufs,
                enable_limiter=enable_limiter,
                enable_spectral_subtract=enable_spectral_subtract,
                enable_comfort_noise=enable_comfort_noise,
                presence_gain_db=presence_gain_db,
                comp_threshold_db=comp_threshold_db,
                comp_ratio=comp_ratio,
                target_lufs=target_lufs,
                spectral_alpha=spectral_alpha,
                spectral_beta=spectral_beta,
                deess_threshold_db=deess_threshold_db,
                deess_ratio=deess_ratio,
            )

        # --------------------------------------------------------
        # OUTPUT
        # --------------------------------------------------------
        st.subheader("🔊 Output")
        st.audio(_audio_bytes(processed, sr), format="audio/wav")

        filename = (
            f"post"
            f"_pres{presence_gain_db:.1f}"
            f"_comp{comp_ratio:.1f}"
            f"_lufs{target_lufs:.0f}.wav"
        )

        st.download_button(
            "Download processed audio",
            _audio_bytes(processed, sr),
            file_name=filename,
            mime="audio/wav",
        )

# ================================================================
# 🔹 Enhancement (SELF-CONTAINED + SLIDERS + CACHED ANALYSIS)
# ================================================================
st.subheader("🎛 Enhancement")

# ---------------------------------------------------------------
# Sidebar controls
# ---------------------------------------------------------------
with st.sidebar:
    st.divider()
    st.header("Input")

    uploaded_file = st.file_uploader("Upload 4-channel WAV", type=["wav"])

    audio_file = st.selectbox(
        "Audio file",
        ["data/raw/example_mixture.wav", "data/raw/mixture.wav"],
    )

    st.subheader("Enhancement Tuning")

    mask_sharpness = st.slider("Mask sharpness", 1.0, 5.0, 2.5, 0.1)
    wiener_floor = st.slider("Wiener floor", 0.01, 0.3, 0.1, 0.01)

    reg = st.slider(
        "Regularization (reg)",
        min_value=1e-8,
        max_value=1e-1,
        value=1e-6,
        step=1e-6,
        format="%.1e",
    )

    st.subheader("Beamforming")

    separation_method = st.selectbox(
        "Method",
        ["mvdr", "das"],
        index=0,
    )

    use_wpe = st.toggle(
        "Use WPE",
        value=False,
    )

    st.markdown("### Target-only controls")

    background_level = st.slider("Background level", 0.0, 0.5, 0.12, 0.01)
    crossfade_ms = st.slider("Crossfade (ms)", 5.0, 100.0, 30.0, 5.0)

    st.markdown("### Spatial rendering")

    ambient_level_db = st.slider("Ambient mix (dB)", -40.0, -5.0, -18.0, 1.0)

# ---------------------------------------------------------------
# Cache invalidation (robust)
# ---------------------------------------------------------------
if (
    "last_method" not in st.session_state
    or st.session_state["last_method"] != separation_method
    or st.session_state.get("last_wpe") != use_wpe
):
    for k in [
        "analysis_done",
        "audio_spatial",
        "sr_spatial",
        "doas",
        "target_idx",
        "target_doa",
        "enhanced_all",
        "asr",
        "ovd_segments",
        "target_ts",
    ]:
        st.session_state.pop(k, None)

st.session_state["last_method"] = separation_method
st.session_state["last_wpe"] = use_wpe

# ---------------------------------------------------------------
# File handling
# ---------------------------------------------------------------
if uploaded_file is not None:
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp.write(uploaded_file.getvalue())
        audio_file = tmp.name

if not Path(audio_file).exists():
    st.error(f"File not found: {audio_file}")
    st.stop()

audio_raw, sr = load_mixture(audio_file)

if audio_raw.ndim != 2 or audio_raw.shape[1] != 4:
    st.error("Expected 4-channel audio")
    st.stop()

# ------------------------------------------------------------
# STEP 1 — ANALYSIS (RUN ONCE)
# ------------------------------------------------------------
if not st.session_state.get("analysis_done", False):

    st.info("Run analysis once. Then tweak parameters interactively.")

    if st.button("🚀 Run Analysis Pipeline", type="primary"):

        with st.spinner("Running analysis pipeline..."):

            audio_spatial, sr_spatial = preprocess_for_spatial(audio_raw, sr)

            doa_result = estimate_doa(audio_spatial, sr_spatial)
            dir_result = estimate_speaker_directions(audio_raw, sr)
            doas = [d["doa"] for d in dir_result["directions"]]

            if not doas:
                st.error("No DoAs detected")
                st.stop()

            target_doa = doas[0]
            target_idx = 0

            enhanced_all = enhance_all_directions(
                audio_spatial,
                sr_spatial,
                doas,
                method=separation_method,
                use_wpe=use_wpe,
                mask_sharpness=mask_sharpness,
                wiener_floor=wiener_floor,
                reg=reg,
            )

            enhanced_ml = []
            for sig in enhanced_all:
                sig_ml, _ = preprocess_enhanced(sig, sr_spatial)
                enhanced_ml.append(sig_ml)

            target_ml = enhanced_ml[target_idx]

            asr = transcribe_detailed(
                target_ml.astype(np.float32),
                16000,
                model_size="medium",
            )

            dialogue = detect_dialogue_pattern(asr["segments"])

            ovd_segments = None
            target_ts = []

            if dialogue["is_dialogue"]:
                ovd_segments = estimate_own_voice_segments(
                    asr["segments"],
                    audio_raw,
                    sr,
                    enhanced_mono=target_ml,
                    enhanced_sr=16000,
                )

                target_ts = [
                    (s["start"], s["end"])
                    for s in ovd_segments
                    if not s.get("is_own_voice") and s.get("text", "").strip()
                ]

        st.session_state.update({
            "audio_spatial": audio_spatial,
            "sr_spatial": sr_spatial,
            "doas": doas,
            "target_doa": target_doa,
            "target_idx": target_idx,
            "enhanced_all": enhanced_all,
            "asr": asr,
            "ovd_segments": ovd_segments,
            "target_ts": target_ts,
            "analysis_done": True,
        })

        st.success("Analysis complete.")
        st.rerun()
    st.stop()

# ------------------------------------------------------------
# STEP 2 — SAFE ACCESS
# ------------------------------------------------------------
required_keys = [
    "audio_spatial",
    "sr_spatial",
    "doas",
    "target_idx",
    "target_doa",
    "enhanced_all",
]

missing = [k for k in required_keys if k not in st.session_state]

if missing:
    st.warning("Please run the analysis pipeline first.")
    st.stop()

audio_spatial = st.session_state["audio_spatial"]
sr_spatial = st.session_state["sr_spatial"]
doas = st.session_state["doas"]
target_idx = st.session_state["target_idx"]
target_doa = st.session_state["target_doa"]
enhanced_all = st.session_state["enhanced_all"]
asr = st.session_state["asr"]
ovd_segments = st.session_state["ovd_segments"]
target_ts = st.session_state["target_ts"]

# ------------------------------------------------------------
# STEP 3 — FAST RENDER
# ------------------------------------------------------------
if st.button("🔄 Regenerate Enhanced Audio", type="primary"):

    with st.spinner("Rendering..."):

        if target_ts:
            enhanced = enhance_target_speaker_only(
                audio_spatial,
                sr_spatial,
                target_doa=target_doa,
                target_timestamps=target_ts,
                doa_list=doas,
                method=separation_method,
                use_wpe=use_wpe,
                mask_sharpness=mask_sharpness,
                wiener_floor=wiener_floor,
                background_level=background_level,
                crossfade_ms=crossfade_ms,
            )
        else:
            enhanced = enhanced_all[target_idx]

        stereo = render_binaural(
            enhanced,
            sr_spatial,
            target_doa=target_doa,
            raw_multichannel=audio_raw,
            ambient_level_db=ambient_level_db,
        )

    st.subheader("🔊 Output")
    st.audio(_audio_bytes(stereo, sr_spatial), format="audio/wav")

    st.subheader("📝 Transcript")

    if ovd_segments:
        for i, seg in enumerate(ovd_segments):
            speaker = "TARGET" if not seg.get("is_own_voice") else "OTHER"

            st.markdown(
                f"**Turn {i+1} — {speaker}** "
                f"({seg['start']:.2f}s → {seg['end']:.2f}s)"
            )
            st.write(seg.get("text", "").strip())
    else:
        st.write(asr["text"])

# ===================================================================
# EXPERIMENT MODE — Preprocessing A/B tests
# ===================================================================

elif mode == "Experiment":
    st.title("Preprocessing Experiment")
    st.markdown(
        "Compare how different preprocessing steps affect transcription and "
        "gender classification on the **raw mixture** (no beamforming). "
        "This isolates the effect of preprocessing alone."
    )

    with st.sidebar:
        st.header("Settings")
        uploaded_file = st.file_uploader(
            "Upload a 4-channel WAV file",
            type=["wav"],
            help="Upload your own multi-channel hearing-aid recording. Must be a 4-channel WAV file.",
        )
        if uploaded_file is not None:
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                tmp.write(uploaded_file.getvalue())
                audio_file = tmp.name
            st.sidebar.success(f"Uploaded: {uploaded_file.name}")
        else:
            audio_file = st.selectbox(
                "Audio file",
                ["data/raw/example_mixture.wav", "data/raw/mixture.wav"],
            )
        whisper_model = st.selectbox(
            "Whisper model",
            ["tiny", "base", "small", "medium", "large"],
            index=3,
            help="'medium' or 'large' recommended for Danish speech.",
        )
        run_asr = st.checkbox("Run transcription (Whisper)", value=True)
        run_gender = st.checkbox("Run gender classification", value=True)

    if not Path(audio_file).exists():
        st.error(f"File not found: `{audio_file}`")
        st.stop()

    audio_raw, sr = load_mixture(audio_file)

    # Validate 4-channel audio
    if audio_raw.ndim != 2 or audio_raw.shape[1] != 4:
        st.error(
            f"Expected a 4-channel WAV file, got shape {audio_raw.shape}. "
            "The file must have exactly 4 channels in order: "
            "[left front, left rear, right front, right rear]."
        )
        st.stop()

    st.sidebar.success(f"Loaded: {audio_raw.shape[1]}ch, {sr} Hz, {audio_raw.shape[0]/sr:.1f}s")

    versions = create_audio_versions(audio_raw, sr)

    if st.button("Run experiment", type="primary"):
        with st.spinner("Running models on all preprocessing versions..."):
            results = run_experiment_on_versions(
                versions, whisper_model=whisper_model,
                run_asr=run_asr, run_gender=run_gender,
            )
        st.session_state["exp_results"] = results
        st.session_state["exp_versions"] = versions
        save_demo_results(results, versions)

    if "exp_results" in st.session_state:
        results = st.session_state["exp_results"]
        versions = st.session_state["exp_versions"]

        df = pd.DataFrame(results)
        display_cols = ["version", "description", "sr", "duration_s", "rms", "peak"]
        if run_asr:
            display_cols += ["asr_language", "asr_segments", "transcript"]
        if run_gender:
            display_cols += ["gender", "gender_confidence", "gender_method", "median_f0"]
        available_cols = [c for c in display_cols if c in df.columns]
        st.dataframe(df[available_cols], use_container_width=True, hide_index=True)

        # Compare two versions
        st.subheader("Compare two versions")
        version_names = list(versions.keys())
        col1, col2 = st.columns(2)
        with col1:
            pick_a = st.selectbox("Version A", version_names, index=0)
        with col2:
            pick_b = st.selectbox("Version B", version_names,
                                  index=min(len(version_names) - 1, 4))

        a_info, b_info = versions[pick_a], versions[pick_b]
        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"**{pick_a}**")
            st.caption(a_info["description"])
            st.audio(_audio_bytes(a_info["audio"], a_info["sr"]), format="audio/wav")
            fig_a = _make_spectrogram_fig(a_info["audio"], a_info["sr"], pick_a)
            st.pyplot(fig_a)
            plt.close(fig_a)
        with col2:
            st.markdown(f"**{pick_b}**")
            st.caption(b_info["description"])
            st.audio(_audio_bytes(b_info["audio"], b_info["sr"]), format="audio/wav")
            fig_b = _make_spectrogram_fig(b_info["audio"], b_info["sr"], pick_b)
            st.pyplot(fig_b)
            plt.close(fig_b)

        csv = df[available_cols].to_csv(index=False)
        st.download_button("Download results CSV", csv, "experiment_results.csv", "text/csv")