import os, json, sys, numpy as np, gradio as gr
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
from huggingface_hub import snapshot_download

PRIVATE = os.environ.get("MODEL_REPO", "")
d = snapshot_download(PRIVATE, token=os.environ.get("HF_TOKEN")) if PRIVATE else os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, d)
import keep_audio_core as K
cfg = json.load(open(os.path.join(d, "config.json"))) if os.path.exists(os.path.join(d, "config.json")) else {}
R = K.Recognizer(os.path.join(d, "audio2b_KEEP_graine0.pt"), os.path.join(d, "audio2b_normalisation.pt"), gain=cfg.get("gain", "peak"))
TAU, MIN_FRAME = 0.95, 10

# Voice-first shell: pure black canvas, one narrow column, no cards, no boxes.
# The result is a stage (huge word + bars), the waveform is the visual centerpiece.
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;700&display=swap');
html, body { background: #000000 !important; }
.gradio-container {
  --body-background-fill: #000000;
  --background-fill-primary: #000000;
  --background-fill-secondary: #000000;
  --block-background-fill: transparent;
  --block-border-color: transparent;
  --block-label-background-fill: transparent;
  --block-label-text-color: #8A8A93;
  --block-title-text-color: #F5F5F7;
  --body-text-color: #F5F5F7;
  --body-text-color-subdued: #8A8A93;
  --button-primary-background-fill: #F5F5F7;
  --button-primary-background-fill-hover: #FFFFFF;
  --button-primary-text-color: #000000;
  --button-secondary-background-fill: transparent;
  --button-secondary-background-fill-hover: #1A1A1E;
  --button-secondary-text-color: #F5F5F7;
  --button-secondary-border-color: rgba(255,255,255,.22);
  --input-background-fill: transparent;
  --input-border-color: rgba(255,255,255,.22);
  --input-text-color: #F5F5F7;
  --input-placeholder-color: #6E6E73;
  max-width: 560px !important; margin: 0 auto !important; padding: 26px 22px 44px !important;
  background: #000000 !important; color: #F5F5F7 !important;
  font-family: 'Space Grotesk', system-ui, sans-serif !important;
}
.gradio-container .block, .gradio-container .form, .gradio-container .wrap,
.gradio-container .panel, .gradio-container .component, .gradio-container .audio,
.gradio-container .plot, .gradio-container .input, .gradio-container .output {
  background: transparent !important; border: none !important;
  box-shadow: none !important; border-radius: 0 !important;
}
.gradio-container label span {
  color: #8A8A93 !important; font-size: 11px !important; font-weight: 500 !important;
  letter-spacing: 2.5px !important; text-transform: uppercase !important;
}
.gradio-container .prose { color: #C7C7CC !important; }
.gradio-container .prose strong { color: #F5F5F7 !important; }
.topbar { display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px; }
.brand { font-weight: 700; font-size: 19px; letter-spacing: -.5px; }
.brand i { color: #E5C15D; font-style: normal; }
.ondevice { font-size: 12px; color: #8A8A93; border: 1px solid rgba(255,255,255,.18); border-radius: 999px; padding: 5px 12px; }
.stage { text-align: center; padding: 34px 0 6px; }
.stage .eyebrow { font-size: 11px; letter-spacing: 3px; text-transform: uppercase; color: #8A8A93; margin-bottom: 16px; }
.stage .word { font-size: 78px; font-weight: 700; letter-spacing: -3px; line-height: 1; margin: 6px 0 0; }
.stage .word.dim { color: #333338; }
.stage .conf { font-size: 15px; color: #E5C15D; margin-top: 14px; }
.stage .top3 { margin-top: 20px; display: flex; flex-direction: column; gap: 8px; }
.stage .top3 .row { display: flex; align-items: center; gap: 12px; font-size: 14px; color: #C7C7CC; }
.stage .top3 .row b { width: 62px; text-align: right; font-weight: 500; color: #F5F5F7; }
.stage .top3 .bar { flex: 1; height: 5px; background: #1C1C1F; border-radius: 99px; overflow: hidden; }
.stage .top3 .bar i { display: block; height: 100%; background: linear-gradient(90deg, #C9A227, #F5D98B); border-radius: 99px; }
.stage .top3 .pct { width: 46px; text-align: left; color: #8A8A93; }
.eq { display: flex; justify-content: center; align-items: flex-end; gap: 5px; height: 46px; margin: 6px 0 10px; }
.eq i { display: block; width: 5px; border-radius: 99px; background: linear-gradient(180deg, #F5D98B, #C9A227); animation: eqpulse 1.15s ease-in-out infinite; transform-origin: bottom; }
@keyframes eqpulse { 0%, 100% { transform: scaleY(.22); } 50% { transform: scaleY(1); } }
.hairline { height: 1px; background: rgba(255,255,255,.10); margin: 30px 0; }
.chips { display: flex; flex-wrap: wrap; gap: 8px; justify-content: center; }
.chip { font-size: 14px; color: #C7C7CC; background: #141416; border-radius: 999px; padding: 9px 15px; }
.quiet { color: #8A8A93 !important; font-size: 13.5px !important; line-height: 1.65 !important; }
.quiet strong { color: #C7C7CC !important; }
.gradio-container button.primary { border-radius: 999px !important; height: 56px !important; font-size: 17px !important; font-weight: 700 !important; }
.gradio-container button:not(.primary) { border-radius: 999px !important; }
.gradio-container .audio button { border-radius: 999px !important; background: #141416 !important; border: 1px solid rgba(255,255,255,.14) !important; }
.gradio-container label[data-testid$="-radio-label"] {
  border-radius: 999px !important; border: 1px solid rgba(255,255,255,.14) !important;
  padding: 7px 15px !important; background: transparent !important; color: #8A8A93 !important;
}
.gradio-container label[data-testid$="-radio-label"].selected { border-color: #E5C15D !important; color: #F5F5F7 !important; }
.gradio-container label[data-testid$="-radio-label"] span { color: inherit !important; background: transparent !important; letter-spacing: 0 !important; text-transform: none !important; font-size: 14px !important; }
.footer { opacity: .45; }
"""

TOPBAR = """<div class="topbar"><div class="brand">KEEP <i>•</i> Listen</div><div class="ondevice">On-device · CPU</div></div>"""

def _eq():
    heights = [22, 34, 18, 40, 28, 46, 26, 36, 20, 30, 42, 24, 32, 18, 38]
    return "".join(f'<i style="height:{h}px;animation-delay:{i * 0.09:.2f}s"></i>' for i, h in enumerate(heights))

def stage_empty():
    return f"""<div class="stage"><div class="eyebrow">KEEP is listening</div><div class="eq">{_eq()}</div><p class="word dim">Say a word</p></div>"""

def stage_result(word, conf, top3):
    rows = "".join(
        f'<div class="row"><b>{w}</b><span class="bar"><i style="width:{p * 100:.0f}%"></i></span><span class="pct">{p:.0%}</span></div>'
        for w, p in top3)
    return (f"""<div class="stage"><div class="eyebrow">KEEP heard</div><p class="word">{word}</p>"""
            f"""<p class="conf">{conf:.0%} confident</p><div class="top3">{rows}</div></div>""")

def glow_plot(P, top, sure):
    t = np.arange(P.shape[0]) * 10
    fig, ax = plt.subplots(figsize=(7, 2.8), facecolor="#000000")
    ax.set_facecolor("#000000")
    COLORS = ["#E5C15D", "#5EC8D8", "#A78BFA"]
    for i, c in zip(top, COLORS):
        ax.plot(t, P[:, i], lw=7, color=c, alpha=0.12, solid_capstyle="round")
        ax.plot(t, P[:, i], lw=2.2, color=c, label=K.WORDS[i], solid_capstyle="round")
    ax.fill_between(t, P[:, top[0]], color="#E5C15D", alpha=0.10)
    ax.axhline(TAU, color="#6E6E73", ls=(0, (2, 4)), lw=1)
    if len(sure):
        td = (sure[0] + MIN_FRAME) * 10
        ax.axvline(td, color="#FFFFFF", ls="--", lw=1, alpha=0.7)
    ax.set_ylim(0, 1.05); ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)
    ax.tick_params(axis="x", colors="#6E6E73", labelsize=9)
    ax.set_xticks([0, 250, 500, 750, 1000])
    ax.set_xlabel("listened time (ms)", color="#6E6E73", fontsize=9)
    leg = ax.legend(loc="upper left", frameon=False, fontsize=10)
    for txt in leg.get_texts():
        txt.set_color("#C7C7CC")
    fig.tight_layout()
    return fig

def listen(audio, gain_choice):
    R.gain = gain_choice
    if audio is None:
        return stage_empty(), gr.update(visible=False), "Record or upload a word first."
    sr, wav = audio
    _, P = R(wav, sr)                                    # P : [101 steps (10 ms), 10 words]
    final = P[-1]; top = final.argsort()[::-1][:3]
    sure = np.where(P[MIN_FRAME:].max(1) >= TAU)[0]
    fig = glow_plot(P, top, sure)
    w = K.WORDS[int(final.argmax())]
    top3 = [(K.WORDS[int(i)], float(final[i])) for i in top]
    stage = stage_result(w, float(final.max()), top3)
    if final.max() < 0.6:
        msg = f"**Not sure** (best guess: “{w}”, {final.max():.0%}). Say one of the 10 words clearly."
    elif len(sure):
        w0 = K.WORDS[int(P[sure[0] + MIN_FRAME].argmax())]
        msg = f"It was already sure of “{w0}” after **{(sure[0] + MIN_FRAME) * 10} ms** of listening."
    else:
        msg = "Recognized at the end of listening."
    return stage, gr.update(value=fig, visible=True), msg

CHIPS = """<div class="chips"><span class="chip">yes</span><span class="chip">no</span><span class="chip">up</span><span class="chip">down</span><span class="chip">left</span><span class="chip">right</span><span class="chip">on</span><span class="chip">off</span><span class="chip">stop</span><span class="chip">go</span></div>"""

with gr.Blocks(title="KEEP Listen", css=CSS) as demo:
    gr.HTML(TOPBAR)
    stage = gr.HTML(stage_empty())
    a = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Your word — about 1 second")
    btn = gr.Button("Listen", variant="primary")
    plot = gr.Plot(show_label=False, visible=False)
    msg = gr.Markdown(elem_classes=["quiet"])
    gain_radio = gr.Radio(["peak", "none"], value=cfg.get("gain", "peak"),
        label="Volume setting (peak = automatic, none = raw)")
    btn.click(listen, [a, gain_radio], [stage, plot, msg])
    gr.HTML('<div class="hairline"></div>')
    gr.Markdown("**Known words.** Say a single word, close to the microphone. The curve shows how its confidence rises while listening.", elem_classes=["quiet"])
    gr.HTML(CHIPS)
    gr.HTML('<div class="hairline"></div>')
    gr.Markdown("**Limits.** It only knows these 10 words: any other sound will be classified as one of them. "
                "Tested at about 97% on never-heard voices (Speech Commands test set); a noisy microphone or a very "
                "different accent can lower that number. It recognizes words, it does not understand speech.",
                elem_classes=["quiet"])
demo.launch()