app.py

"""Gradio demo and inference API for the acoustic-5 MeowContext Lab baseline."""

from __future__ import annotations

import argparse
import base64
import io
import subprocess
import sys
import tempfile
import wave
from pathlib import Path

import numpy as np

_SRC = Path(__file__).resolve().parent / "src"
if _SRC.is_dir() and str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))

from meowcontext_lab.data import DEMO_MODEL_PATH, FEATURE_COLUMNS  # noqa: E402
from meowcontext_lab.models import load_demo_model, predict_from_features  # noqa: E402

MAX_DURATION_SEC = 30.0
RECOMMENDED_MAX_SEC = 10.0
MIN_RMS = 0.005
TARGET_SAMPLE_RATE = 8000


def _load_wav_array(path: str | Path) -> tuple[np.ndarray, int]:
    with wave.open(str(path), "rb") as wav:
        sample_rate = wav.getframerate()
        channels = wav.getnchannels()
        frames = wav.getnframes()
        raw = wav.readframes(frames)
        sample_width = wav.getsampwidth()

    if sample_width == 1:
        audio = np.frombuffer(raw, dtype=np.uint8).astype(np.float32)
        audio = (audio - 128) / 128
    elif sample_width == 2:
        audio = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768
    else:
        raise ValueError("Only 8-bit and 16-bit WAV files are supported.")

    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)
    if len(audio) == 0:
        raise ValueError("Audio file is empty.")
    return audio, sample_rate


def _ffmpeg_to_wav(source: str | Path, destination: str | Path) -> None:
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(source),
        "-vn",
        "-acodec",
        "pcm_s16le",
        "-ar",
        str(TARGET_SAMPLE_RATE),
        "-ac",
        "1",
        str(destination),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        stderr = (result.stderr or "").strip()
        if "does not contain any stream" in stderr or "Output file is empty" in stderr:
            raise ValueError("Video has no audio track.")
        raise ValueError("Could not extract audio from the uploaded file.")


def load_audio_array(path: str | Path) -> tuple[np.ndarray, int]:
    """Load mono float audio from WAV or convert via ffmpeg."""

    path = Path(path)
    if path.suffix.lower() == ".wav":
        audio, sample_rate = _load_wav_array(path)
    else:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
            wav_path = Path(handle.name)
        try:
            _ffmpeg_to_wav(path, wav_path)
            audio, sample_rate = _load_wav_array(wav_path)
        finally:
            wav_path.unlink(missing_ok=True)

    if sample_rate != TARGET_SAMPLE_RATE and len(audio) > 1:
        duration = len(audio) / sample_rate
        target_len = max(1, int(round(duration * TARGET_SAMPLE_RATE)))
        x_old = np.linspace(0, duration, num=len(audio), endpoint=False)
        x_new = np.linspace(0, duration, num=target_len, endpoint=False)
        audio = np.interp(x_new, x_old, audio).astype(np.float32)
        sample_rate = TARGET_SAMPLE_RATE
    return audio, sample_rate


def acoustic5_from_array(audio: np.ndarray, sample_rate: int) -> dict[str, float]:
    duration = len(audio) / sample_rate
    rms = float(np.sqrt(np.mean(np.square(audio))))
    peak = float(np.max(np.abs(audio)))
    zcr = float(np.mean(np.abs(np.diff(np.signbit(audio))))) if len(audio) > 1 else 0.0
    spectrum = np.abs(np.fft.rfft(audio))
    freqs = np.fft.rfftfreq(len(audio), d=1 / sample_rate)
    centroid = float(np.sum(freqs * spectrum) / max(np.sum(spectrum), 1e-12))
    return {
        "duration_sec": float(duration),
        "rms_energy": rms,
        "peak_abs_amplitude": peak,
        "zero_crossing_rate": zcr,
        "spectral_centroid_hz": centroid,
    }


def acoustic5_from_wav(path: str | Path) -> dict[str, float]:
    audio, sample_rate = load_audio_array(path)
    return acoustic5_from_array(audio, sample_rate)


def _png_data_uri(fig) -> str:
    import matplotlib.pyplot as plt

    buffer = io.BytesIO()
    fig.savefig(buffer, format="png", dpi=120, bbox_inches="tight", facecolor="#0f1419")
    plt.close(fig)
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def waveform_image(audio: np.ndarray, sample_rate: int):
    import matplotlib.pyplot as plt

    seconds = np.arange(len(audio)) / sample_rate
    fig, ax = plt.subplots(figsize=(6.5, 2.2), facecolor="#0f1419")
    ax.set_facecolor("#0f1419")
    ax.plot(seconds, audio, color="#7c9cff", linewidth=0.9)
    ax.set_xlabel("Time (s)", color="#cbd5e1", fontsize=9)
    ax.set_ylabel("Amplitude", color="#cbd5e1", fontsize=9)
    ax.tick_params(colors="#94a3b8", labelsize=8)
    ax.set_title("Waveform", color="#e2e8f0", fontsize=10, pad=8)
    for spine in ax.spines.values():
        spine.set_color("#334155")
    fig.tight_layout()
    return _png_data_uri(fig)


def spectrogram_image(audio: np.ndarray, sample_rate: int):
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(6.5, 2.6), facecolor="#0f1419")
    ax.set_facecolor("#0f1419")
    nfft = min(512, max(64, len(audio) // 8))
    spec = ax.specgram(
        audio,
        NFFT=nfft,
        Fs=sample_rate,
        noverlap=nfft // 2,
        cmap="magma",
    )
    ax.set_xlabel("Time (s)", color="#cbd5e1", fontsize=9)
    ax.set_ylabel("Frequency (Hz)", color="#cbd5e1", fontsize=9)
    ax.tick_params(colors="#94a3b8", labelsize=8)
    ax.set_title("Spectrogram", color="#e2e8f0", fontsize=10, pad=8)
    cbar = fig.colorbar(spec[3], ax=ax, fraction=0.046, pad=0.04)
    cbar.ax.tick_params(colors="#94a3b8", labelsize=7)
    cbar.set_label("Power (dB)", color="#cbd5e1", fontsize=8)
    for spine in ax.spines.values():
        spine.set_color("#334155")
    fig.tight_layout()
    return _png_data_uri(fig)


def _collect_warnings(features: dict[str, float]) -> list[str]:
    warnings: list[str] = []
    duration = features["duration_sec"]
    if duration > MAX_DURATION_SEC:
        warnings.append(f"Clip is longer than {MAX_DURATION_SEC:.0f}s and may be unreliable.")
    elif duration > RECOMMENDED_MAX_SEC:
        warnings.append(f"Clip exceeds the recommended {RECOMMENDED_MAX_SEC:.0f}s window.")
    if features["rms_energy"] < MIN_RMS:
        warnings.append("Clip is very quiet; prediction may be unreliable.")
    if features["peak_abs_amplitude"] < 0.01:
        warnings.append("Very low peak amplitude detected.")
    return warnings


def build_prediction_response(path: str | Path) -> dict:
    """Run acoustic-5 inference and return structured JSON for the public website."""

    if not path:
        raise ValueError("No audio file provided.")

    if not DEMO_MODEL_PATH.exists():
        raise FileNotFoundError(
            f"{DEMO_MODEL_PATH} not found. Run `python scripts/train_demo_model.py` first."
        )

    audio, sample_rate = load_audio_array(path)
    features = acoustic5_from_array(audio, sample_rate)
    bundle = load_demo_model(DEMO_MODEL_PATH)
    prediction = predict_from_features(bundle, features)
    probabilities = {str(k): float(v) for k, v in prediction.probabilities.items()}
    confidence = float(max(probabilities.values()))

    return {
        "predicted_context": prediction.label,
        "confidence": confidence,
        "probabilities": probabilities,
        "warnings": _collect_warnings(features),
        "waveform_image": waveform_image(audio, sample_rate),
        "spectrogram_image": spectrogram_image(audio, sample_rate),
        "features": {key: float(features[key]) for key in FEATURE_COLUMNS},
    }


def predict_audio_api(audio_path: str | None) -> dict:
    """Named API endpoint: /predict_audio"""

    if not audio_path:
        raise ValueError("No audio detected in the upload.")
    return build_prediction_response(audio_path)


def predict_video_api(video_path: str | None) -> dict:
    """Named API endpoint: /predict_video"""

    if not video_path:
        raise ValueError("No video file provided.")
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
        wav_path = Path(handle.name)
    try:
        _ffmpeg_to_wav(video_path, wav_path)
        return build_prediction_response(wav_path)
    finally:
        wav_path.unlink(missing_ok=True)


CONTEXT_LABELS = {
    "brushing": "Brushing",
    "isolation_unfamiliar_environment": "Isolation (unfamiliar place)",
    "waiting_for_food": "Waiting for food",
}


def _confidence_html(probabilities: dict[str, float]) -> str:
    rows: list[str] = []
    for label, value in sorted(probabilities.items(), key=lambda item: item[1], reverse=True):
        name = CONTEXT_LABELS.get(label, label.replace("_", " "))
        percent = int(round(float(value) * 100))
        rows.append(
            f'<div style="margin-bottom:12px">'
            f'<div style="display:flex;justify-content:space-between;color:#cbd5e1;font-size:14px">'
            f"<span>{name}</span><span>{percent}%</span></div>"
            f'<div style="background:#1e293b;border-radius:999px;height:8px;margin-top:6px">'
            f'<div style="width:{max(2, percent)}%;background:linear-gradient(90deg,#6366f1,#a78bfa);'
            f'height:8px;border-radius:999px"></div></div></div>'
        )
    return "".join(rows)


def _image_html(data_uri: str, alt: str) -> str:
    if not data_uri:
        return ""
    return (
        f'<img src="{data_uri}" alt="{alt}" '
        f'style="width:100%;border-radius:12px;margin-top:8px;display:block" />'
    )


def _friendly_result(path: str | Path) -> tuple[str, str, str, str, str]:
    response = build_prediction_response(path)
    label = CONTEXT_LABELS.get(
        response["predicted_context"],
        response["predicted_context"].replace("_", " "),
    )
    confidence = int(round(float(response["confidence"]) * 100))
    headline = f"## Your cat sounds like: **{label}**"
    summary = f"Best match: **{label}** ({confidence}% confidence)"
    warnings = (
        "\n\n".join(f"⚠️ {warning}" for warning in response["warnings"])
        if response["warnings"]
        else ""
    )
    visuals = _image_html(response["waveform_image"], "Waveform") + _image_html(
        response["spectrogram_image"], "Spectrogram"
    )
    if warnings:
        visuals += f"\n\n{warnings}"
    return headline, summary, _confidence_html(response["probabilities"]), visuals, (
        "This is a benchmark demo, not a real-world cat interpretation system."
    )


def predict_audio_ui(audio_path: str | None) -> tuple[str, str, str, str, str]:
    if not audio_path:
        raise ValueError("Record or upload a meow first.")
    return _friendly_result(audio_path)


def predict_video_ui(video_path: str | None) -> tuple[str, str, str, str, str]:
    if not video_path:
        raise ValueError("Upload or record a short cat video first.")
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle:
        wav_path = Path(handle.name)
    try:
        _ffmpeg_to_wav(video_path, wav_path)
        return _friendly_result(wav_path)
    finally:
        wav_path.unlink(missing_ok=True)


def create_demo():
    import gradio as gr

    with gr.Blocks(title="MeowContext Lab", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# MeowContext Lab")
        gr.Markdown("Record or upload a cat meow — see what situation it matches.")
        gr.Markdown(
            "This demo predicts one of three eliciting recording contexts from a tiny public "
            "dataset. It is **not** a cat translator, emotion detector, pain detector, welfare "
            "tool, veterinary tool, or diagnostic system."
        )
        gr.Markdown("Tip: 2–10 second clips work best.")

        result_headline = gr.Markdown()
        result_summary = gr.Markdown()
        result_bars = gr.HTML()
        result_visuals = gr.HTML()
        result_note = gr.Markdown()

        with gr.Tab("Record audio"):
            mic = gr.Audio(sources=["microphone"], type="filepath", label="Record a meow")
            gr.Button("See result", variant="primary").click(
                predict_audio_ui,
                inputs=[mic],
                outputs=[result_headline, result_summary, result_bars, result_visuals, result_note],
            )

        with gr.Tab("Upload audio"):
            audio_file = gr.Audio(sources=["upload"], type="filepath", label="Upload audio")
            gr.Button("See result", variant="primary").click(
                predict_audio_ui,
                inputs=[audio_file],
                outputs=[result_headline, result_summary, result_bars, result_visuals, result_note],
            )

        with gr.Tab("Upload / record video"):
            video_file = gr.Video(sources=["upload", "webcam"], label="Cat video (audio only)")
            gr.Button("See result", variant="primary").click(
                predict_video_ui,
                inputs=[video_file],
                outputs=[result_headline, result_summary, result_bars, result_visuals, result_note],
            )

        gr.Markdown(
            "Uploaded or recorded files are used only for this prediction and are not used to "
            "train the model."
        )

        # Hidden API hooks for the public Vercel website (not shown to visitors).
        with gr.Row(visible=False):
            api_audio = gr.Audio(type="filepath")
            api_audio_out = gr.JSON()
            gr.Button("api_audio").click(
                predict_audio_api,
                inputs=[api_audio],
                outputs=[api_audio_out],
                api_name="predict_audio",
            )
            api_video = gr.Video()
            api_video_out = gr.JSON()
            gr.Button("api_video").click(
                predict_video_api,
                inputs=[api_video],
                outputs=[api_video_out],
                api_name="predict_video",
            )
    return demo


demo = create_demo()


def smoke_test() -> None:
    if not DEMO_MODEL_PATH.exists():
        raise FileNotFoundError(
            f"{DEMO_MODEL_PATH} not found. Run `python scripts/train_demo_model.py` first."
        )
    features = dict(
        zip(
            FEATURE_COLUMNS,
            [1.4, 0.12, 0.36, 0.08, 1200.0],
            strict=True,
        )
    )
    bundle = load_demo_model(DEMO_MODEL_PATH)
    prediction = predict_from_features(bundle, features)
    print(f"Smoke prediction: {prediction.label}")


def _write_tiny_wav(path: Path) -> None:
    sample_rate = 8000
    t = np.linspace(0, 0.5, sample_rate // 2, endpoint=False)
    signal = 0.2 * np.sin(2 * np.pi * 440 * t)
    pcm = (signal * 32767).astype("<i2")
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(pcm.tobytes())


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", nargs="*", help="Use `smoke test` for a CLI smoke test.")
    parser.add_argument("--share", action="store_true", help="Create a public Gradio share URL.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command in (["smoke"], ["smoke", "test"], ["smoke-test"]):
        smoke_test()
        with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
            _write_tiny_wav(Path(handle.name))
            response = build_prediction_response(handle.name)
        print(f"Smoke API context: {response['predicted_context']}")
        print(f"Smoke confidence: {response['confidence']:.3f}")
        return
    demo.launch(share=args.share)


if __name__ == "__main__":
    main()