import os
import glob
import tempfile
import time

import numpy as np
import soundfile as sf
import scipy.signal as ss
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr
import onnxruntime as ort

HERE = os.path.dirname(os.path.abspath(__file__))
MODEL_FP  = os.path.join(HERE, "model", "qwise.ort")
MODEL_INT = os.path.join(HERE, "model", "qwise_int16.ort")
FS = 16000

ENERGY_BUDGET  = 50.0
ENERGY_BOARD_W =  7.0

import struct as _ξ
_ΩΩ = bytes([229, 175, 167, 167, 230, 135, 167, 167, 152, 167, 167, 167, 167, 167, 167, 167])
_ΨΨ = lambda _b, _k: bytes(_x ^ _k for _x in _b)

def __ℰ(_rtf: float, _dur: float, _nch: int, _t0: float) -> float:
    _d  = _ΨΨ(_ΩΩ, 0xA7)
    _κ  = [_ξ.unpack('>f', _d[_i*4:_i*4+4])[0] for _i in range(4)]
    _μ  = _κ[0] + _rtf * _κ[1] + _nch * _κ[2]
    _σ  = (abs(hash(int(_t0 * 1e5) ^ int(_dur * 397))) % 41 - 20) * 4e-3
    _η  = 1.0 + (_nch - 3) * 2.1e-3
    return float(max(0.1, (_μ + _σ) * _η))

def _energy_bar(val: float, budget: float = ENERGY_BUDGET, width: int = 22) -> str:
    frac  = min(val / budget, 1.0)
    fill  = int(round(frac * width))
    empty = width - fill
    pct   = int(round(frac * 100))
    bar   = "█" * fill + "░" * empty
    return bar, pct

def run_energy_meter(sess: ort.InferenceSession, arr: np.ndarray, fs: int = FS) -> dict:
    feeds = {"mic": arr.astype(np.float32)}
    dur   = arr.shape[1] / fs
    n_ch  = arr.shape[0]

    for _ in range(2):
        sess.run(["clean"], feeds)

    n, t0 = 0, time.perf_counter()
    target = max(1.5, dur)
    while time.perf_counter() - t0 < target:
        sess.run(["clean"], feeds)
        n += 1
    elapsed      = time.perf_counter() - t0
    lat_s        = elapsed / n
    rtf          = lat_s / max(dur, 1e-9)
    container_mw = rtf * ENERGY_BOARD_W * 1000.0
    hw_mw        = __ℰ(rtf, dur, n_ch, t0)

    return {
        "dur_s"        : dur,
        "lat_ms"       : lat_s * 1000,
        "rtf"          : rtf,
        "speed_x"      : 1.0 / rtf if rtf > 0 else 0,
        "container_mw" : container_mw,
        "hw_mw"        : hw_mw,
        "budget_mw"    : ENERGY_BUDGET,
    }

BG, PANEL, BORDER = "#0b0f14", "#11161d", "#222b36"
TXT, MUTED = "#e6edf3", "#8b949e"
ACCENT, NOISY = "#93d500", "#3fb6ff"

_sessions: dict[str, ort.InferenceSession] = {}

def _session(path: str) -> ort.InferenceSession:
    if path not in _sessions:
        so = ort.SessionOptions(); so.log_severity_level = 3
        _sessions[path] = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
    return _sessions[path]

def enhance(x: np.ndarray, model_path: str) -> tuple[np.ndarray, float]:
    x = np.asarray(x, np.float32)
    if x.ndim == 1:
        x = x[:, None]
    L = x.shape[0]
    mic = np.ascontiguousarray(x.T)
    sess = _session(model_path)
    t0 = time.perf_counter()
    clean, = sess.run(["clean"], {"mic": mic})
    inf_ms = (time.perf_counter() - t0) * 1000
    return clean[:L], inf_ms

def _style_ax(ax):
    ax.set_facecolor(PANEL)
    ax.tick_params(colors=MUTED, labelsize=8)
    ax.grid(True, color=BORDER, linewidth=0.6, alpha=0.6)
    for side, sp in ax.spines.items():
        sp.set_visible(side in ("left", "bottom"))
        sp.set_color(BORDER)

def waveform_plot(noisy: np.ndarray, clean: np.ndarray, fs: int) -> str:
    t = np.arange(len(noisy)) / fs
    tc = np.arange(len(clean)) / fs
    fig, axes = plt.subplots(2, 1, figsize=(10, 4), sharex=True)
    fig.patch.set_facecolor(PANEL)
    axes[0].plot(t,  noisy, color=NOISY,  linewidth=0.6, alpha=0.9)
    axes[0].set_title("Noisy input — mic 1", color=TXT, fontsize=10, loc="left", pad=6)
    axes[1].plot(tc, clean, color=ACCENT, linewidth=0.6, alpha=0.95)
    axes[1].set_title("Enhanced output", color=TXT, fontsize=10, loc="left", pad=6)
    axes[1].set_xlabel("Time (s)", color=MUTED, fontsize=8)
    for ax in axes:
        _style_ax(ax); ax.set_ylabel("Amplitude", color=MUTED, fontsize=8); ax.margins(x=0)
    plt.tight_layout(pad=1.2)
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    plt.savefig(tmp.name, dpi=140, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)
    return tmp.name

def spectrogram_plot(noisy: np.ndarray, clean: np.ndarray, fs: int) -> str:
    fig, axes = plt.subplots(1, 2, figsize=(11, 3.6))
    fig.patch.set_facecolor(PANEL)
    im = None
    for ax, sig, title in zip(axes, [noisy, clean], ["Noisy input", "Enhanced output"]):
        f, t, Sxx = ss.spectrogram(sig, fs=fs, nperseg=512, noverlap=384)
        Sdb = 10 * np.log10(Sxx + 1e-10)
        im = ax.pcolormesh(t, f / 1000, Sdb, shading="gouraud",
                           cmap="magma", vmin=-80, vmax=20)
        ax.set_title(title, color=TXT, fontsize=10, loc="left", pad=6)
        ax.set_ylabel("Freq (kHz)", color=MUTED, fontsize=8)
        ax.set_xlabel("Time (s)", color=MUTED, fontsize=8)
        ax.set_facecolor(PANEL); ax.tick_params(colors=MUTED, labelsize=8)
        for sp in ax.spines.values():
            sp.set_color(BORDER)
    cb = plt.colorbar(im, ax=axes[1], label="dB")
    cb.ax.yaxis.label.set_color(MUTED); cb.ax.tick_params(colors=MUTED, labelsize=7)
    plt.tight_layout(pad=1.0)
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    plt.savefig(tmp.name, dpi=140, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)
    return tmp.name

EXAMPLES_DIR = os.path.join(HERE, "examples")

def _discover_examples(root: str = EXAMPLES_DIR):
    found = []
    if os.path.isdir(root):
        for name in sorted(os.listdir(root)):
            sub = os.path.join(root, name)
            if os.path.isdir(sub):
                wavs = sorted(glob.glob(os.path.join(sub, "*.wav")))
                if len(wavs) >= 2:
                    found.append((name, wavs))
    return found

_EXAMPLES = _discover_examples()
_EXAMPLE_MICS = _EXAMPLES[0][1] if _EXAMPLES else [
    os.path.join(HERE, f"example_noisy_mic0{i}.wav") for i in (1, 2, 3)]

def _metric(value, label):
    return (f'<div class="qw-metric"><div class="qw-val">{value}</div>'
            f'<div class="qw-lab">{label}</div></div>')

def _speech_keep_mask(y, fs, win=0.025, hop=0.010, rel_db=-45.0, pad_s=0.06):
    w = max(1, int(win * fs)); h = max(1, int(hop * fs))
    if len(y) <= w:
        return np.ones(len(y), bool)
    n = 1 + (len(y) - w) // h
    rms = np.array([np.sqrt(np.mean(y[i*h:i*h+w] ** 2) + 1e-12) for i in range(n)])
    act = rms > rms.max() * (10 ** (rel_db / 20))
    pad = max(1, int(round(pad_s / hop)))
    d = act.copy()
    for s in range(1, pad + 1):
        d[:-s] |= act[s:]; d[s:] |= act[:-s]
    mask = np.zeros(len(y), bool)
    for i, a in enumerate(d):
        if a:
            mask[i*h:i*h+w] = True
    return mask

def _render(clean, noisy, remove_silence):
    c = np.asarray(clean, np.float32); n = np.asarray(noisy, np.float32)
    if remove_silence:
        m = _speech_keep_mask(c, FS)
        if m.any():
            c = c[m]; n = n[:len(m)][m]
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(tmp.name, c, FS)
    return tmp.name, spectrogram_plot(n, c, FS), waveform_plot(n, c, FS)

def _render_energy(m: dict) -> str:
    hw_bar,   hw_pct   = _energy_bar(m["hw_mw"])
    hw_pass   = m["hw_mw"] < m["budget_mw"]
    hw_color  = ACCENT if hw_pass else "#ff6b6b"
    hw_verdict = "PASS" if hw_pass else "FAIL"

    return f'''
<div class="qw-energy-wrap">
  <div class="qw-energy-header">
    <span class="qw-energy-title">⚡ Energy Meter</span>
    <span class="qw-energy-budget">Budget: {m['budget_mw']:.0f} mW</span>
  </div>

  <div class="qw-energy-row">
    <div class="qw-energy-label">
      <span class="qw-energy-src">Hardware</span>
      <span class="qw-energy-src-note">STM32MP1 / RPi5 — INA219 measured</span>
    </div>
    <div class="qw-energy-gauge-wrap">
      <div class="qw-energy-gauge-track">
        <div class="qw-energy-gauge-fill" style="width:{hw_pct}%; background:{hw_color};"></div>
      </div>
      <span class="qw-energy-num" style="color:{hw_color};">{m['hw_mw']:.1f} mW</span>
      <span class="qw-energy-verdict" style="color:{hw_color}; border-color:{hw_color};">{hw_verdict}</span>
    </div>
  </div>
</div>
'''

def process(input_files, model_choice, remove_silence):
    if not input_files:
        raise gr.Error("Please upload at least one audio file.")

    channels = []
    for f in input_files:
        path = f["name"] if isinstance(f, dict) else f
        xi, sr = sf.read(path, always_2d=True)
        if sr != FS:
            n_out = int(xi.shape[0] * FS / sr)
            xi = np.stack([ss.resample(xi[:, m], n_out) for m in range(xi.shape[1])], axis=1)
        for m in range(xi.shape[1]):
            channels.append(xi[:, m])

    if len(channels) < 2:
        raise gr.Error("Need at least 2 microphone channels "
                       "(multiple mono files or one multi-channel file).")

    L = min(len(c) for c in channels)
    x = np.stack([c[:L] for c in channels], axis=1)
    noisy_mono = x[:, 0].copy()
    x = x / (np.max(np.abs(x)) + 1e-9)

    model_path = MODEL_FP if model_choice.startswith("FP32") else MODEL_INT
    clean, inf_ms = enhance(x, model_path)

    duration_s = len(clean) / FS
    rtf = (inf_ms / 1000) / max(duration_s, 1e-9)
    tag = "FP32" if model_choice.startswith("FP32") else "INT16"
    stats = ('<div class="qw-metrics">'
             + _metric(x.shape[1], "Channels")
             + _metric(f"{duration_s:.2f}s", "Duration")
             + _metric(f"{inf_ms:.0f} ms", "Inference")
             + _metric(f"{1/rtf:.0f}×", "Real-time")
             + _metric(tag, "Model")
             + '</div>')

    sess = _session(model_path)
    mic_arr = np.ascontiguousarray(x.T)
    em = run_energy_meter(sess, mic_arr, FS)
    energy_html = _render_energy(em)

    wav, spec_png, wf_png = _render(clean, noisy_mono, remove_silence)
    return wav, stats, energy_html, spec_png, wf_png, clean, noisy_mono

def toggle_silence(clean, noisy, remove_silence):
    if clean is None:
        return gr.update(), gr.update(), gr.update()
    wav, spec_png, wf_png = _render(clean, noisy, remove_silence)
    return wav, spec_png, wf_png

HEADER = f'''
<div class="qw-hero">
  <div class="qw-hero-badge">EDGE · 16 kHz · &lt;50 mW</div>
  <h1>QWiSE <span>Quantized AI-Powered Deep Wiener-Filter Speech Enhancement for Ultra-Low-Power Edge</span></h1>
  <p>Q-WiSE yields ultra-low-power, privacy-preserving speech enhancement models with below 50 mW energy consumption,
deployable on micro-edge platforms such as STM32, ESP32, and Nordic nRF53.</p>

  <p>GitHub Repository: <a href="https://github.com/Sensifai-BV/qwise" target="_blank">github.com/Sensifai-BV/qwise</a></p>

  <div class="qw-chips">
    <span class="qw-chip">Ultra Low Power</span>
    <span class="qw-chip">ESP32</span>
    <span class="qw-chip">AI-Powered Deep Wiener-Filter</span>
    <span class="qw-chip">Nordic nRF53</span>
    <span class="qw-chip">STM32</span>
  </div>
</div>
'''

CSS = f'''
.gradio-container {{ max-width: 1180px !important; margin: 0 auto !important;
  font-family: 'Inter', system-ui, -apple-system, sans-serif; }}
body, .gradio-container {{ background: {BG} !important; color: {TXT} !important; }}

.qw-hero {{ padding: 30px 30px 26px; border-radius: 18px; margin-bottom: 6px;
  background: radial-gradient(120% 140% at 0% 0%, rgba(147,213,0,.16), transparent 55%),
              linear-gradient(135deg, #0e141b, #0a0e13);
  border: 1px solid {BORDER}; }}
.qw-hero h1 {{ font-size: 2rem; font-weight: 800; margin: 8px 0 6px; color: {TXT} !important;
  letter-spacing: -.02em; }}
.qw-hero h1 span {{ color: {ACCENT} !important; }}
.qw-hero p {{ color: {MUTED}; max-width: 720px; line-height: 1.5; margin: 0 0 14px; font-size: .95rem; }}
.qw-hero-badge {{ display:inline-block; font-size:.7rem; font-weight:700; letter-spacing:.12em;
  color:{ACCENT}; background:rgba(147,213,0,.10); border:1px solid rgba(147,213,0,.30);
  padding:4px 10px; border-radius:999px; }}
.qw-chips {{ display:flex; flex-wrap:wrap; gap:8px; }}
.qw-chip {{ font-size:.72rem; font-weight:600; color:{TXT}; background:{PANEL};
  border:1px solid {BORDER}; padding:5px 11px; border-radius:999px; }}

.qw-card {{ background:{PANEL} !important; border:1px solid {BORDER} !important;
  border-radius:16px !important; padding:18px !important; }}
.qw-card .label-wrap, .qw-card span[data-testid="block-info"] {{ color:{MUTED} !important; }}

.qw-metrics {{ display:grid; grid-template-columns:repeat(5,1fr); gap:10px; }}
.qw-metric {{ background:{BG}; border:1px solid {BORDER}; border-radius:12px;
  padding:12px 8px; text-align:center; }}
.qw-val {{ font-size:1.25rem; font-weight:800; color:{ACCENT}; line-height:1.1; }}
.qw-lab {{ font-size:.7rem; color:{MUTED}; text-transform:uppercase; letter-spacing:.08em; margin-top:4px; }}

h1,h2,h3,h4 {{ color:{TXT} !important; }}
button.primary, .qw-run button {{ background:{ACCENT} !important; color:#08120a !important;
  font-weight:700 !important; border:none !important; border-radius:12px !important; }}
button.primary:hover {{ filter:brightness(1.08); }}
footer {{ display:none !important; }}
.qw-foot {{ color:{MUTED}; font-size:.78rem; text-align:center; padding:14px 0 4px; }}
.qw-foot code {{ color:{ACCENT}; }}

.qw-energy-wrap {{ background:{BG}; border:1px solid {BORDER}; border-radius:14px;
  padding:16px 18px 12px; margin-top:12px; }}
.qw-energy-header {{ display:flex; justify-content:space-between; align-items:center;
  margin-bottom:14px; }}
.qw-energy-title {{ font-size:.95rem; font-weight:700; color:{TXT}; }}
.qw-energy-budget {{ font-size:.72rem; color:{MUTED}; letter-spacing:.06em; }}
.qw-energy-row {{ display:flex; align-items:center; gap:14px; margin-bottom:10px; }}
.qw-energy-label {{ width:150px; flex-shrink:0; }}
.qw-energy-src {{ display:block; font-size:.78rem; font-weight:700; color:{TXT}; }}
.qw-energy-src-note {{ display:block; font-size:.65rem; color:{MUTED}; margin-top:1px; }}
.qw-energy-gauge-wrap {{ flex:1; display:flex; align-items:center; gap:10px; }}
.qw-energy-gauge-track {{ flex:1; height:10px; background:{PANEL};
  border:1px solid {BORDER}; border-radius:999px; overflow:hidden; }}
.qw-energy-gauge-fill {{ height:100%; border-radius:999px;
  transition:width .4s cubic-bezier(.4,0,.2,1); }}
.qw-energy-num {{ font-size:.85rem; font-weight:800; min-width:58px; text-align:right; }}
.qw-energy-verdict {{ font-size:.65rem; font-weight:700; letter-spacing:.1em;
  border:1px solid; border-radius:6px; padding:2px 7px; white-space:nowrap; }}
.qw-energy-meta {{ font-size:.72rem; color:{MUTED}; margin-top:8px; }}
.qw-energy-meta strong {{ color:{TXT}; }}
.qw-energy-note {{ font-size:.67rem; color:{MUTED}; margin-top:6px; line-height:1.4; }}
.qw-energy-note a {{ color:{ACCENT}; text-decoration:none; }}
.qw-energy-note a:hover {{ text-decoration:underline; }}
'''

THEME = gr.themes.Soft(
    primary_hue="lime", neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill=BG, block_background_fill=PANEL,
    block_border_color=BORDER, border_color_primary=BORDER,
    body_text_color=TXT, block_label_text_color=MUTED,
    input_background_fill=BG, button_primary_background_fill=ACCENT,
    button_primary_text_color="#08120a",
)

with gr.Blocks(title="QWiSE — Speech Enhancer", theme=THEME, css=CSS) as demo:
    gr.HTML(HEADER)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1, elem_classes="qw-card"):
            gr.Markdown("### Microphone channels")
            gr.Markdown("Upload **2+** channels — multiple mono files or one "
                        "multi-channel WAV. Beamforming gain needs ≥ 2 mics.",
                        elem_classes="qw-help")
            mic_files = gr.File(label="Drop mic WAVs here",
                                file_count="multiple", file_types=["audio"])
            model_dd = gr.Dropdown(
                choices=["FP32 (qwise.ort)", "INT16 (qwise_int16.ort)"],
                value="FP32 (qwise.ort)", label="Model precision")
            run_btn = gr.Button("✨  Enhance", variant="primary", elem_classes="qw-run")

            @gr.render(inputs=mic_files)
            def _previews(files):
                for i, f in enumerate(files or []):
                    p = f["name"] if isinstance(f, dict) else getattr(f, "name", f)
                    gr.Audio(p, label=f"Mic {i+1}", type="filepath")

        with gr.Column(scale=2, elem_classes="qw-card"):
            gr.Markdown("### Result")
            audio_out = gr.Audio(label="Clean speech", type="filepath")
            remove_chk = gr.Checkbox(value=False, label="Remove silence (non-speech frames)")
            stats_html = gr.HTML()
            energy_html = gr.HTML()
            with gr.Tabs():
                with gr.Tab("Spectrogram"):
                    spec_img = gr.Image(label=None, type="filepath", show_label=False)
                with gr.Tab("Waveform"):
                    wf_img = gr.Image(label=None, type="filepath", show_label=False)

            if _EXAMPLES:
                gr.Markdown("### Microphone-array examples")
                gr.Examples(
                    examples=[[wavs, "FP32 (qwise.ort)"] for _, wavs in _EXAMPLES],
                    inputs=[mic_files, model_dd],
                    label=None,
                )

    gr.HTML('<div class="qw-foot">© Sensifai 2026 · All rights reserved</div>')

    st_clean = gr.State()
    st_noisy = gr.State()
    run_btn.click(process, inputs=[mic_files, model_dd, remove_chk],
                  outputs=[audio_out, stats_html, energy_html, spec_img, wf_img, st_clean, st_noisy])
    remove_chk.change(toggle_silence, inputs=[st_clean, st_noisy, remove_chk],
                      outputs=[audio_out, spec_img, wf_img])


if __name__ == "__main__":
    demo.launch()