app.py

'''Whistle Coach - audio-first Hugging Face Space.

The app uses MIT AST audio classification for whistle confidence and
librosa.pyin for pitch tracking. Camera guidance is intentionally visual-only.
'''

from __future__ import annotations

import base64
import io
import json
import math
from pathlib import Path
from typing import Any

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
from scipy.signal import resample_poly
from transformers import pipeline

try:
    import spaces
except Exception:
    class _SpacesFallback:
        @staticmethod
        def GPU(*decorator_args, **decorator_kwargs):
            if (
                decorator_args
                and callable(decorator_args[0])
                and len(decorator_args) == 1
                and not decorator_kwargs
            ):
                return decorator_args[0]

            def decorator(function):
                return function

            return decorator

    spaces = _SpacesFallback()

MODEL_ID = "MIT/ast-finetuned-audioset-10-10-0.4593"
AST_PARAMETERS = "86.6M"
AST_DEVICE = "cpu"
TARGET_SR = 16_000
MIN_AUDIO_SECONDS = 0.25
MAX_ANALYSIS_SECONDS = 4.0

BASE_DIR = Path(__file__).resolve().parent
MELODY_DIR = BASE_DIR / "outputs" / "melodies"
MELODY_DIR.mkdir(parents=True, exist_ok=True)

MODEL_LOAD_ERROR: str | None = None
try:
    audio_classifier = pipeline(
        "audio-classification",
        model=MODEL_ID,
        device=-1,
    )
except Exception as exc:
    audio_classifier = None
    MODEL_LOAD_ERROR = f"{type(exc).__name__}: {exc}"

NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
STATE_ORDER = ["no_sound", "breath_noise", "tiny_whistle", "stable_pitch", "melody_ready"]
STATE_LABELS = {
    "no_sound": "No sound",
    "breath_noise": "Breath noise",
    "tiny_whistle": "Tiny whistle",
    "stable_pitch": "Stable pitch",
    "melody_ready": "Melody",
}
GARDEN_STAGES = {
    "no_sound": ("soil", "Soil", "Rest, breathe, and start softly."),
    "breath_noise": ("wind", "Wind", "Air is moving. Make the opening smaller."),
    "tiny_whistle": ("sprout", "Sprout", "A tiny tone appeared. Freeze that shape."),
    "stable_pitch": ("flower", "Flower", "Hold the tone and keep it gentle."),
    "melody_ready": ("melody", "Melody flower", "You can shape notes into a melody."),
}


def clamp(value: float, minimum: float = 0.0, maximum: float = 100.0) -> float:
    return float(min(maximum, max(minimum, value)))


def hz_to_note(hz: float | None) -> str:
    if hz is None or not np.isfinite(hz) or hz <= 0:
        return "--"
    midi = int(round(69 + 12 * math.log2(float(hz) / 440.0)))
    octave = midi // 12 - 1
    return f"{NOTE_NAMES[midi % 12]}{octave}"


def hz_to_midi(hz: float) -> int:
    return int(round(69 + 12 * math.log2(float(hz) / 440.0)))


def midi_to_hz(midi: int) -> float:
    return 440.0 * (2 ** ((midi - 69) / 12))


def normalize_audio_array(data: np.ndarray) -> np.ndarray:
    y = np.asarray(data)
    if y.ndim > 1:
        y = np.mean(y, axis=1)
    if np.issubdtype(y.dtype, np.integer):
        peak = float(np.iinfo(y.dtype).max)
        y = y.astype(np.float32) / max(peak, 1.0)
    else:
        y = y.astype(np.float32)
        max_abs = float(np.max(np.abs(y))) if y.size else 0.0
        if max_abs > 1.5:
            y = y / max_abs
    y = np.nan_to_num(y, nan=0.0, posinf=0.0, neginf=0.0)
    return np.clip(y, -1.0, 1.0).astype(np.float32)


def load_audio_input(audio_input: Any) -> tuple[int, np.ndarray]:
    '''Accept Gradio Audio numpy tuples, file paths, or path dictionaries.'''
    if audio_input is None:
        return TARGET_SR, np.zeros(0, dtype=np.float32)

    if isinstance(audio_input, tuple) and len(audio_input) == 2:
        sample_rate, data = audio_input
        return int(sample_rate), normalize_audio_array(np.asarray(data))

    if isinstance(audio_input, dict):
        path = audio_input.get("path") or audio_input.get("name")
        if path:
            data, sample_rate = sf.read(path, always_2d=False)
            return int(sample_rate), normalize_audio_array(np.asarray(data))

    if isinstance(audio_input, (str, Path)):
        data, sample_rate = sf.read(str(audio_input), always_2d=False)
        return int(sample_rate), normalize_audio_array(np.asarray(data))

    return TARGET_SR, np.zeros(0, dtype=np.float32)


def resample_to_target(sample_rate: int, y: np.ndarray) -> np.ndarray:
    if y.size == 0:
        return y.astype(np.float32)
    if sample_rate == TARGET_SR:
        return y.astype(np.float32)
    gcd = math.gcd(int(sample_rate), TARGET_SR)
    up = TARGET_SR // gcd
    down = int(sample_rate) // gcd
    return resample_poly(y, up, down).astype(np.float32)


def latest_window(y: np.ndarray, seconds: float = MAX_ANALYSIS_SECONDS) -> np.ndarray:
    max_samples = int(TARGET_SR * seconds)
    if y.size <= max_samples:
        return y
    return y[-max_samples:]


def classify_with_ast(y16: np.ndarray) -> tuple[float, list[dict[str, Any]], str | None]:
    '''Run the real HF AST classifier and return whistle confidence.'''
    if audio_classifier is None:
        return 0.0, [], MODEL_LOAD_ERROR or "Audio model is not loaded."
    if y16.size < int(TARGET_SR * MIN_AUDIO_SECONDS):
        return 0.0, [], None

    try:
        predictions = audio_classifier(
            {"array": y16.astype(np.float32), "sampling_rate": TARGET_SR},
            top_k=8,
        )
    except TypeError:
        predictions = audio_classifier(
            {"array": y16.astype(np.float32), "sampling_rate": TARGET_SR}
        )
    except Exception as exc:
        return 0.0, [], f"AST inference failed: {type(exc).__name__}: {exc}"

    if isinstance(predictions, dict):
        predictions = [predictions]
    if predictions and isinstance(predictions[0], list):
        predictions = predictions[0]

    labels: list[dict[str, Any]] = []
    whistle_score = 0.0
    top_non_whistle = 0.0
    for item in predictions[:8]:
        label = str(item.get("label", ""))
        score = float(item.get("score", 0.0) or 0.0)
        labels.append({"label": label, "score": round(score, 4)})
        normalized = label.lower().replace("_", " ").replace("-", " ")
        if "whistl" in normalized or "whistle" in normalized:
            whistle_score = max(whistle_score, score)
        else:
            top_non_whistle = max(top_non_whistle, score)

    if whistle_score > 0:
        confidence = whistle_score * 100.0
    else:
        confidence = clamp((1.0 - top_non_whistle) * 12.0, 0.0, 12.0)
    return round(confidence, 2), labels, None


def longest_stable_duration(f0: np.ndarray, voiced: np.ndarray, hop_length: int) -> float:
    longest = 0.0
    current: list[float] = []

    def finish_run(values: list[float]) -> float:
        if len(values) < 3:
            return 0.0
        arr = np.asarray(values, dtype=np.float64)
        center = float(np.nanmedian(arr))
        if center <= 0:
            return 0.0
        cents = 1200.0 * np.log2(arr / center)
        if float(np.nanstd(cents)) > 95.0:
            return 0.0
        return len(values) * hop_length / TARGET_SR

    for hz, is_voiced in zip(f0, voiced):
        if bool(is_voiced) and np.isfinite(hz) and hz > 0:
            current.append(float(hz))
        else:
            longest = max(longest, finish_run(current))
            current = []
    longest = max(longest, finish_run(current))
    return round(float(longest), 2)


def analyze_pitch(y16: np.ndarray) -> dict[str, Any]:
    hop_length = 256
    frame_length = 2048
    if y16.size < frame_length:
        return {
            "pitch_detected": False,
            "pitch_note": "--",
            "mean_pitch_hz": None,
            "stability_score": 0,
            "stable_duration": 0.0,
            "pitch_contour": [],
            "voiced_ratio": 0.0,
            "pitch_std_cents": None,
            "pitch_movement_semitones": 0.0,
        }

    try:
        f0, voiced_flag, voiced_prob = librosa.pyin(
            y16,
            fmin=librosa.note_to_hz("C5"),
            fmax=librosa.note_to_hz("C8"),
            sr=TARGET_SR,
            frame_length=frame_length,
            hop_length=hop_length,
        )
    except Exception:
        return {
            "pitch_detected": False,
            "pitch_note": "--",
            "mean_pitch_hz": None,
            "stability_score": 0,
            "stable_duration": 0.0,
            "pitch_contour": [],
            "voiced_ratio": 0.0,
            "pitch_std_cents": None,
            "pitch_movement_semitones": 0.0,
        }

    voiced = np.asarray(voiced_flag, dtype=bool) & np.isfinite(f0)
    valid_f0 = np.asarray(f0)[voiced]
    voiced_ratio = float(np.mean(voiced)) if len(voiced) else 0.0
    pitch_detected = bool(valid_f0.size >= 3 and voiced_ratio >= 0.06)

    mean_pitch: float | None = None
    pitch_std_cents: float | None = None
    pitch_movement = 0.0
    if valid_f0.size:
        mean_pitch = float(np.nanmedian(valid_f0))
        if mean_pitch > 0 and valid_f0.size > 1:
            cents = 1200.0 * np.log2(valid_f0 / mean_pitch)
            pitch_std_cents = float(np.nanstd(cents))
            pitch_movement = float(np.nanmax(cents) - np.nanmin(cents)) / 100.0

    std_score = 0.0
    if pitch_std_cents is not None:
        std_score = clamp(100.0 - (pitch_std_cents * 1.35), 0.0, 100.0)
    stability_score = clamp((voiced_ratio * 62.0) + (std_score * 0.38))
    stable_duration = longest_stable_duration(np.asarray(f0), voiced, hop_length)

    times = librosa.frames_to_time(np.arange(len(f0)), sr=TARGET_SR, hop_length=hop_length)
    contour: list[dict[str, Any]] = []
    valid_indices = np.where(np.isfinite(f0) & voiced)[0]
    if valid_indices.size:
        step = max(1, int(math.ceil(valid_indices.size / 90)))
        for idx in valid_indices[::step]:
            hz = float(f0[idx])
            contour.append({
                "time_ms": int(round(float(times[idx]) * 1000)),
                "hz": round(hz, 2),
                "note": hz_to_note(hz),
            })

    return {
        "pitch_detected": pitch_detected,
        "pitch_note": hz_to_note(mean_pitch),
        "mean_pitch_hz": round(mean_pitch, 2) if mean_pitch else None,
        "stability_score": int(round(stability_score)) if pitch_detected else 0,
        "stable_duration": stable_duration if pitch_detected else 0.0,
        "pitch_contour": contour,
        "voiced_ratio": round(voiced_ratio, 3),
        "pitch_std_cents": round(pitch_std_cents, 2) if pitch_std_cents is not None else None,
        "pitch_movement_semitones": round(pitch_movement, 2),
    }


def analyze_airflow(y16: np.ndarray) -> dict[str, float]:
    if y16.size == 0:
        return {"rms": 0.0, "db": -120.0, "airflow_score": 0.0, "noise_score": 0.0}
    rms = float(np.sqrt(np.mean(np.square(y16))))
    db = float(20.0 * np.log10(rms + 1e-9))
    airflow_score = clamp((db + 55.0) / 32.0 * 100.0)
    try:
        flatness = float(np.nanmean(librosa.feature.spectral_flatness(y=y16, n_fft=1024, hop_length=256)))
        zcr = float(np.nanmean(librosa.feature.zero_crossing_rate(y16, frame_length=1024, hop_length=256)))
    except Exception:
        flatness = 0.0
        zcr = 0.0
    noise_score = clamp((flatness * 180.0) + (zcr * 110.0))
    return {
        "rms": round(rms, 5),
        "db": round(db, 2),
        "airflow_score": int(round(airflow_score)),
        "noise_score": round(noise_score, 2),
    }


def choose_state_and_feedback(result: dict[str, Any]) -> tuple[str, str, str]:
    airflow = float(result["airflow_score"])
    confidence = float(result["whistle_confidence"])
    pitch_detected = bool(result["pitch_detected"])
    stable_duration = float(result["stable_duration"] or 0.0)
    noise_score = float(result.get("noise_score", 0.0) or 0.0)
    movement = float(result.get("pitch_movement_semitones", 0.0) or 0.0)

    if pitch_detected and stable_duration >= 1.0 and movement >= 2.0:
        return (
            "melody_ready",
            "Nice - you are changing notes. Try making a melody.",
            "Move the pitch slowly, like stepping between two nearby notes.",
        )
    if pitch_detected and stable_duration >= 1.0:
        return (
            "stable_pitch",
            "Great! Hold this tone longer.",
            "Stay still for one more second before trying to move the pitch.",
        )
    if pitch_detected:
        return (
            "tiny_whistle",
            "Tiny whistle found. Freeze this mouth shape.",
            "Use less air now; keep the same lip opening and hold it steady.",
        )
    if airflow < 10:
        return (
            "no_sound",
            "Blow a little more, but stay gentle.",
            "Keep your lips relaxed, then send a small stream of air through the center.",
        )
    if confidence >= 25:
        return (
            "tiny_whistle",
            "You are close. Make the air stream narrower.",
            "Round the lips a little more and soften the airflow.",
        )
    if airflow >= 35 or noise_score >= 18:
        return (
            "breath_noise",
            "You are producing air noise. Make the lip opening smaller and soften the airflow.",
            "Try a quiet 'yuh yuh yuh' breath with a smaller opening.",
        )
    return (
        "no_sound",
        "Blow a little more, but stay gentle.",
        "Stay relaxed. The first goal is a soft, steady stream of air.",
    )


def analyze_audio(audio_input: Any) -> dict[str, Any]:
    '''Analyze a Gradio Audio input using AST and librosa.pyin.'''
    sample_rate, y = load_audio_input(audio_input)
    y16 = latest_window(resample_to_target(sample_rate, y))
    duration = y16.size / TARGET_SR if TARGET_SR else 0.0

    airflow = analyze_airflow(y16)
    whistle_confidence, ast_labels, ast_error = classify_with_ast(y16)
    pitch = analyze_pitch(y16)

    result: dict[str, Any] = {
        "airflow_score": int(airflow["airflow_score"]),
        "whistle_confidence": int(round(whistle_confidence)),
        "pitch_detected": bool(pitch["pitch_detected"]),
        "pitch_note": pitch["pitch_note"],
        "mean_pitch_hz": pitch["mean_pitch_hz"],
        "stability_score": int(pitch["stability_score"]),
        "stable_duration": float(pitch["stable_duration"]),
        "state": "no_sound",
        "coach_feedback": "Blow a little more, but stay gentle.",
        "next_tip": "Start with a soft, narrow stream of air.",
        "pitch_contour": pitch["pitch_contour"],
        "audio_seconds": round(duration, 2),
        "rms": airflow["rms"],
        "db": airflow["db"],
        "noise_score": airflow["noise_score"],
        "voiced_ratio": pitch["voiced_ratio"],
        "pitch_std_cents": pitch["pitch_std_cents"],
        "pitch_movement_semitones": pitch["pitch_movement_semitones"],
        "ast_top_labels": ast_labels,
        "ast_error": ast_error,
        "model_stack": {
            "audio_model": f"{MODEL_ID} ({AST_PARAMETERS}, {AST_DEVICE})",
            "pitch_tracking": "librosa.pyin C5-C8",
            "visual_assistant": "MediaPipe visible mouth guidance only; no tongue detection",
        },
    }

    state, feedback, next_tip = choose_state_and_feedback(result)
    result["state"] = state
    result["coach_feedback"] = feedback
    result["next_tip"] = next_tip
    return result


def contour_to_note_sequence(contour: list[dict[str, Any]]) -> list[dict[str, Any]]:
    notes: list[dict[str, Any]] = []
    current_midi: int | None = None
    start_ms = 0
    last_ms = 0
    for sample in contour:
        hz = float(sample.get("hz", sample.get("frequency_hz", 0)) or 0)
        t_ms = int(sample.get("time_ms", sample.get("t", 0)) or 0)
        if hz < 120:
            continue
        midi = hz_to_midi(hz)
        if current_midi is None:
            current_midi = midi
            start_ms = t_ms
        elif abs(midi - current_midi) > 1:
            notes.append({"midi": current_midi, "duration_ms": max(180, last_ms - start_ms)})
            current_midi = midi
            start_ms = t_ms
        last_ms = t_ms
    if current_midi is not None:
        notes.append({"midi": current_midi, "duration_ms": max(260, last_ms - start_ms)})
    return notes[:12]


def generate_melody_from_contour(contour: list[dict[str, Any]]) -> tuple[str | None, str]:
    notes = contour_to_note_sequence(contour)
    if not notes:
        return None, ""
    sample_rate = 22_050
    samples: list[np.ndarray] = []
    for note in notes:
        hz = midi_to_hz(int(note["midi"]))
        duration = min(0.75, max(0.18, int(note["duration_ms"]) / 1000.0))
        t = np.linspace(0.0, duration, int(sample_rate * duration), endpoint=False)
        tone = np.sin(2 * np.pi * hz * t) + 0.25 * np.sin(2 * np.pi * hz * 2 * t)
        fade = min(220, max(12, len(tone) // 8))
        envelope = np.ones_like(tone)
        envelope[:fade] = np.linspace(0.0, 1.0, fade)
        envelope[-fade:] = np.linspace(1.0, 0.0, fade)
        samples.append((tone * envelope * 0.23).astype(np.float32))
        samples.append(np.zeros(int(sample_rate * 0.045), dtype=np.float32))
    wav = np.concatenate(samples) if samples else np.zeros(1, dtype=np.float32)
    output_path = MELODY_DIR / "whistle_melody.wav"
    sf.write(output_path, wav, sample_rate)
    names = [hz_to_note(midi_to_hz(int(note["midi"]))) for note in notes]
    return str(output_path), " - ".join(names)


FRONTEND_DIR = BASE_DIR / "frontend"
HTML_PATH = FRONTEND_DIR / "index.html"
STYLE_PATH = FRONTEND_DIR / "style.css"
SCRIPT_PATHS = [
    FRONTEND_DIR / "policy.js",
    FRONTEND_DIR / "landmarks.js",
    FRONTEND_DIR / "audio_features.js",
    FRONTEND_DIR / "face_features.js",
    FRONTEND_DIR / "overlay_renderer.js",
    FRONTEND_DIR / "whistle_coach.js",
]


def load_frontend_html() -> str:
    return HTML_PATH.read_text(encoding="utf-8")


def load_frontend_css() -> str:
    return STYLE_PATH.read_text(encoding="utf-8") if STYLE_PATH.exists() else ""


def load_js_bundle() -> str:
    parts: list[str] = []
    for script_path in SCRIPT_PATHS:
        if script_path.exists():
            parts.append(script_path.read_text(encoding="utf-8"))
    return "\n\n".join(parts)


@spaces.GPU(duration=1)
def zerogpu_startup_probe() -> str:
    """Tiny hidden hook so ZeroGPU Spaces start without charging live audio calls."""
    return "ready"


def analyze_audio_window(audio_payload_json: str) -> dict[str, Any]:
    """Decode a browser-captured WAV window and run the real AST/librosa stack."""
    if not audio_payload_json:
        result = empty_result()
        result["ast_error"] = MODEL_LOAD_ERROR
        return result

    try:
        payload = json.loads(audio_payload_json)
        encoded = payload.get("data_base64") or ""
        audio_bytes = base64.b64decode(encoded)
        data, sample_rate = sf.read(io.BytesIO(audio_bytes), always_2d=False)
    except Exception as exc:
        result = empty_result()
        result["ast_error"] = f"Audio decode failed: {type(exc).__name__}: {exc}"
        result["coach_feedback"] = "The microphone window could not be decoded. Try again."
        result["next_tip"] = "Keep the browser microphone permission enabled."
        return result

    result = analyze_audio((int(sample_rate), data))
    result["source"] = "browser_live_window"
    return result


def empty_result() -> dict[str, Any]:
    return {
        "airflow_score": 0,
        "whistle_confidence": 0,
        "pitch_detected": False,
        "pitch_note": "--",
        "mean_pitch_hz": None,
        "stability_score": 0,
        "stable_duration": 0.0,
        "state": "no_sound",
        "coach_feedback": "Ready when you are.",
        "next_tip": "The coach listens to each audio window and updates feedback.",
        "pitch_contour": [],
        "rms": 0,
        "db": -120,
        "voiced_ratio": 0,
        "pitch_std_cents": None,
        "pitch_movement_semitones": 0,
        "ast_top_labels": [],
        "ast_error": MODEL_LOAD_ERROR,
        "model_stack": {
            "audio_model": f"{MODEL_ID} ({AST_PARAMETERS}, {AST_DEVICE})",
            "pitch_tracking": "librosa.pyin C5-C8",
            "visual_assistant": "MediaPipe visible mouth guidance only; no tongue detection",
        },
    }


def coach_reason(practice_state_json: str, history_json: str = "[]") -> dict[str, Any]:
    """Small rule-based policy kept for the old frontend trace panel."""
    try:
        practice_state = json.loads(practice_state_json or "{}")
    except json.JSONDecodeError:
        practice_state = {}
    state_name = practice_state.get("state", "idle")
    audio = practice_state.get("audio", {}) if isinstance(practice_state, dict) else {}
    airflow = float(audio.get("airflow_score", 0) or 0)
    stability = float(audio.get("pitch_stability_score", 0) or 0)

    if stability >= 0.58:
        message = "Great! Hold this tone longer."
        drill = "Hold a steady tone"
        success = True
    elif airflow >= 0.55:
        message = "You are producing air noise. Make the lip opening smaller and soften the airflow."
        drill = "Make the air stream narrower"
        success = False
    elif airflow >= 0.25:
        message = "You are close. Make the air stream narrower."
        drill = "Gentle airflow"
        success = False
    else:
        message = "Blow a little more, but stay gentle."
        drill = "Gentle airflow"
        success = False

    return {
        "coach_message": message,
        "next_drill": drill,
        "success_trigger": success,
        "model_source": "local rule fallback; AST/librosa panel uses live audio model",
        "agent_trace": [
            {"step": "Observe", "detail": f"Browser practice state: {state_name}."},
            {"step": "Listen", "detail": "Backend AST/librosa analyzes one-second audio windows."},
            {"step": "Coach", "detail": drill},
        ],
    }


def generate_melody(pitch_contour_json: str) -> dict[str, Any]:
    try:
        contour = json.loads(pitch_contour_json or "[]")
    except json.JSONDecodeError:
        contour = []
    path, sequence = generate_melody_from_contour(contour if isinstance(contour, list) else [])
    return {"path": path, "sequence": sequence, "ready": bool(path)}


with gr.Blocks(
    title="Whistle Coach",
    fill_width=True,
) as demo:
    gr.HTML(load_frontend_html(), container=False, padding=False)

    zerogpu_trigger = gr.Button(visible=False)
    zerogpu_status = gr.Textbox(visible=False)
    zerogpu_trigger.click(
        zerogpu_startup_probe,
        outputs=zerogpu_status,
        api_name="zerogpu_startup_probe",
        show_progress="hidden",
    )

    audio_payload = gr.Textbox(visible=False)
    audio_result = gr.JSON(visible=False)
    audio_payload.submit(
        analyze_audio_window,
        inputs=audio_payload,
        outputs=audio_result,
        api_name="analyze_audio_window",
        show_progress="hidden",
    )

    practice_state = gr.Textbox(visible=False)
    coach_history = gr.Textbox(visible=False)
    coach_output = gr.JSON(visible=False)
    practice_state.submit(
        coach_reason,
        inputs=[practice_state, coach_history],
        outputs=coach_output,
        api_name="coach_reason",
        show_progress="hidden",
    )

    melody_contour = gr.Textbox(visible=False)
    melody_output = gr.JSON(visible=False)
    melody_contour.submit(
        generate_melody,
        inputs=melody_contour,
        outputs=melody_output,
        api_name="generate_melody",
        show_progress="hidden",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        theme=gr.themes.Soft(),
        css=load_frontend_css(),
        js=load_js_bundle(),
    )