frontend/layout.py

from __future__ import annotations
import base64
import json
import re
import time
from pathlib import Path
import gradio as gr
import soundfile as sf
import config
from backend.presets import list_presets
from frontend import live
LOG_HEADERS = ['t_s', 'model', 'decision', 'probability', 'latency_ms']

def _load_sample_manifest() -> list[dict]:
    manifest_path = config.SAMPLE_CLIPS_DIR / 'manifest.json'
    if not manifest_path.exists():
        return []
    return json.loads(manifest_path.read_text())
BUCKET_LABEL = {'english': 'English', 'hindi': 'Hindi', 'hinglish': 'Hinglish'}
BUCKET_ORDER = ['english', 'hindi', 'hinglish']

def _sample_clip_choices() -> list[tuple[str, str]]:
    manifest = _load_sample_manifest()
    by_bucket: dict[str, list[dict]] = {b: [] for b in BUCKET_ORDER}
    for entry in manifest:
        by_bucket.setdefault(entry.get('bucket', 'other'), []).append(entry)
    choices = []
    for bucket in BUCKET_ORDER:
        for i, entry in enumerate(by_bucket.get(bucket, []), start=1):
            label = f'{BUCKET_LABEL.get(bucket, bucket.title())} {i} - {entry['duration_seconds']}s'
            choices.append((label, entry['filename']))
    for bucket, entries in by_bucket.items():
        if bucket in BUCKET_ORDER:
            continue
        for i, entry in enumerate(entries, start=1):
            choices.append((f'{bucket.title()} {i} - {entry['duration_seconds']}s', entry['filename']))
    return choices

def _all_public_presets() -> list[dict]:
    return live.all_public_presets()
FUTURE_PRESET_LABELS = ['Easy Turn (unavailable)']

def _display_presets() -> list[dict]:
    presets = live.all_public_presets()
    seen = {p['label'] for p in presets}
    for label in FUTURE_PRESET_LABELS:
        if label not in seen:
            match = next((p for p in list_presets() if p['label'] == label), None)
            if match:
                presets.append(match)
    return presets

def _preset_display_order() -> list[str]:
    return [live.DISPLAY_NAME.get(p['label'], p['label']).replace(' (unavailable)', '') for p in _display_presets()]

def _default_active_display_names() -> list[str]:
    target = 'Whisper-Tiny + Mean-Pool + Linear (trained)'
    name = live.DISPLAY_NAME.get(target, target)
    names = _preset_display_order()
    return [name] if name in names else [names[0]] if names else []

def _on_preset_change(*values: bool) -> tuple[list[str], dict, dict]:
    names = _preset_display_order()
    active = [n for n, v in zip(names, values) if v]
    ac, se = _param_relevance(active)
    return (active, gr.update(interactive=ac), gr.update(interactive=se))

def _param_relevance(active_display_names: list[str]) -> tuple[bool, bool]:
    display_to_internal = {v: k for k, v in live.DISPLAY_NAME.items()}
    labels = [display_to_internal[n] for n in active_display_names or [] if n in display_to_internal]
    acoustic_relevant = any((live.uses_acoustic_weight(label) for label in labels))
    semantic_relevant = any((live.uses_semantic_temperature(label) for label in labels))
    return (acoustic_relevant, semantic_relevant)

def _load_doc_markdown() -> str:
    root = Path(__file__).resolve().parent.parent
    doc_path = root / 'docs' / 'Turn Detection - Aman.md'
    if not doc_path.exists():
        return '_Documentation is not included in this deployment._'
    text = doc_path.read_text(encoding='utf-8')
    for a, b in [('\\&quot;', '"'), ('\\&amp;', '&'), ('\\&apos;', "'"), ('\\&lt;', '<'), ('\\&gt;', '>')]:
        text = text.replace(a, b)

    def repl(m: re.Match) -> str:
        alt, src = (m.group(1), m.group(2).strip())
        if src.startswith(('http', 'data:')):
            return m.group(0)
        cand = (doc_path.parent / src).resolve()
        if cand.exists():
            b64 = base64.b64encode(cand.read_bytes()).decode('ascii')
            return f'![{alt}](data:image/png;base64,{b64})'
        return m.group(0)
    return re.sub('!\\[([^\\]]*)\\]\\(([^)]+)\\)', repl, text)
CUSTOM_CSS = '\n.wrap-row { flex-wrap: wrap; gap: 10px; }\n.wrap-row > * { flex: 1 1 220px; min-width: 170px; }\n'

def build_app() -> gr.Blocks:
    with gr.Blocks(title='Turn Detection - Live Dashboard') as demo:
        gr.Markdown('# Turn Detection - Live Dashboard\nSpeak, or replay a clip, and watch how different models judge whether the speaker is **done talking** vs. **still going** - plotted directly against the waveform as audio arrives. Open the **Documentation** tab for the full write-up, methodology and results.')
        with gr.Tabs():
            with gr.Tab('Live Dashboard', id='live'):
                session_state = gr.State(live.new_session_state)
                active_state = gr.State(_default_active_display_names())
                gr.Markdown('### Models to compare  *(greyed = not built yet / future work)*')
                with gr.Row(elem_classes='wrap-row'):
                    default_active = _default_active_display_names()
                    preset_checkboxes = []
                    for preset in _display_presets():
                        label = preset['label']
                        disp = live.DISPLAY_NAME.get(label, label).replace(' (unavailable)', '')
                        available = bool(preset.get('available'))
                        cb_label = disp if available else f'{disp}  (future)'
                        preset_checkboxes.append(gr.Checkbox(label=cb_label, value=disp in default_active, interactive=available))
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown('### Audio')
                        mic = gr.Audio(sources=['microphone'], streaming=True, type='numpy', label='Record')
                        gr.Markdown('*Recordings are saved privately to improve the model - never played back or shown to others.*')
                        clear_btn = gr.Button('Clear / reset')
                        gr.Markdown('**...or replay a clip in real time** (paced to its real duration, not dumped in at once) - 10 real English, 10 Hindi, 10 Hinglish')
                        replay_clip_dropdown = gr.Dropdown(choices=_sample_clip_choices(), value=None, label='Sample clip')
                        replay_upload = gr.Audio(sources=['upload'], type='filepath', label='...or upload a recording')
                        replay_btn = gr.Button('Replay in real time')
                    with gr.Column(scale=2):
                        gr.Markdown('### Waveform + live probability, on one timeline')
                        chart = gr.Plot(value=live.render_chart(live.new_session_state(), [], live.DEFAULT_DECISION_THRESHOLD), label=None)
                with gr.Row(elem_classes='wrap-row'):
                    threshold_slider = gr.Slider(0.0, 1.0, value=live.DEFAULT_DECISION_THRESHOLD, step=0.01, label='Decision threshold - probability above this = "complete"')
                    cadence_slider = gr.Slider(0, 5000, value=live.DEFAULT_CADENCE_MS, step=100, label='Update cadence (ms) - how often each model re-checks (slower models may still lag behind this)')
                    smoothing_slider = gr.Slider(0.0, 0.9, value=0.5, step=0.05, label="Smoothing - damps short swings on the chart (raw points still shown faintly; doesn't change what's logged)")
                    acoustic_weight_slider = gr.Slider(0.0, 1.0, value=0.6, label='Acoustic weight - audio tone vs. sentence grammar (fusion models only)', interactive=False)
                    temperature_slider = gr.Slider(0.0, 1.0, value=0.2, label='Semantic temperature - how deterministic the language judgment is (LLM-based models only)', interactive=False)
                with gr.Accordion('History (this session)', open=False):
                    log_table = gr.Dataframe(headers=LOG_HEADERS, value=[], label=None)
                for cb in preset_checkboxes:
                    cb.change(_on_preset_change, inputs=preset_checkboxes, outputs=[active_state, acoustic_weight_slider, temperature_slider])

                def replay_clip(clip_filename, uploaded_path, active_display_names, acoustic_weight, temperature, threshold, cadence_ms, smoothing):
                    path = uploaded_path or (str(config.SAMPLE_CLIPS_DIR / clip_filename) if clip_filename else None)
                    if not path:
                        yield (live.new_session_state(), live.render_chart(live.new_session_state(), [], threshold, smoothing), [])
                        return
                    audio, sr = sf.read(path, dtype='float32')
                    if audio.ndim > 1:
                        audio = audio.mean(axis=1)
                    state = live.new_session_state()
                    for chunk in live.chunk_audio(audio, sr, chunk_seconds=1.0):
                        chunk_duration_s = len(chunk) / sr
                        state, fig, log_rows = live.process_chunk(state, (sr, chunk), active_display_names, acoustic_weight, temperature, threshold, cadence_ms, smoothing)
                        yield (state, fig, log_rows)
                        time.sleep(chunk_duration_s)
                replay_btn.click(replay_clip, inputs=[replay_clip_dropdown, replay_upload, active_state, acoustic_weight_slider, temperature_slider, threshold_slider, cadence_slider, smoothing_slider], outputs=[session_state, chart, log_table])
                mic.stream(fn=live.process_chunk, inputs=[session_state, mic, active_state, acoustic_weight_slider, temperature_slider, threshold_slider, cadence_slider, smoothing_slider], outputs=[session_state, chart, log_table], stream_every=1.0, time_limit=None)
                clear_btn.click(live.clear_session, inputs=session_state, outputs=[session_state, chart, log_table])
            with gr.Tab('Documentation', id='docs'):
                gr.HTML('<style>.doc-md { max-height: 78vh; overflow: auto; padding-right: 16px; }.doc-md img { max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 6px; margin: 8px 0; }</style>')
                gr.Markdown(_load_doc_markdown(), elem_classes=['doc-md'])
    return demo