app.py

import io
import math
import time
import tempfile
from datetime import datetime

import cv2
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.signal import butter, filtfilt, find_peaks


# ----------------------------
# Signal processing utilities
# ----------------------------
def _butter_bandpass(lowcut: float, highcut: float, fs: float, order: int = 4):
    nyq = 0.5 * fs
    low = max(0.0001, lowcut / nyq)
    high = min(0.9999, highcut / nyq)
    from scipy.signal import butter as _butter
    b, a = _butter(order, [low, high], btype="band")
    return b, a


def _bandpass_filter(signal: np.ndarray, lowcut=0.7, highcut=4.0, fs=30.0, order=4):
    b, a = _butter_bandpass(lowcut, highcut, fs, order=order)
    return filtfilt(b, a, signal)  # zero-phase


def video_to_mean_green_trace(video_path: str, fallback_fps: float = 30.0):
    """Read video and return (mean_green_trace, fps, duration_seconds)."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError("Cannot open video")

    fps = cap.get(cv2.CAP_PROP_FPS)
    if not fps or fps <= 0 or fps > 240:
        fps = fallback_fps

    trace = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        trace.append(float(np.mean(rgb[:, :, 1])))  # mean of green channel
    cap.release()

    x = np.asarray(trace, dtype=np.float32)
    duration = float(len(x) / fps) if fps > 0 else 0.0
    return x, float(fps), duration


def estimate_bpm_and_quality(trace: np.ndarray, fps: float):
    """Return (bpm, filtered_trace, quality_score_0_100) or (None, filtered, None) if unreliable."""
    if trace is None or len(trace) < 10 or fps <= 0:
        return None, None, None

    # normalize
    sig = (trace - np.mean(trace)) / (np.std(trace) + 1e-8)

    # band-pass in typical HR range: ~0.7–4.0 Hz (42–240 bpm)
    try:
        filtered = _bandpass_filter(sig, 0.7, 4.0, fs=fps, order=4)
    except Exception:
        filtered = sig - np.mean(sig)

    # peak detection with a sane min distance (~150 bpm upper bound)
    min_distance = max(1, int(0.4 * fps))
    peaks, _ = find_peaks(filtered, distance=min_distance)
    if len(peaks) < 2:
        return None, filtered, None

    # bpm from mean RR interval
    intervals = np.diff(peaks) / fps
    if np.any(intervals <= 0):
        return None, filtered, None
    bpm = 60.0 / np.mean(intervals)

    # quality: simple SNR-like score
    noise = sig - filtered
    p_sig = float(np.mean(np.square(filtered)))
    p_noise = float(np.mean(np.square(noise))) + 1e-10
    snr_db = 10.0 * math.log10(p_sig / p_noise)
    quality = max(0.0, min(100.0, (snr_db + 10.0) * 5.0))  # maps approx -10..10 dB -> 0..100

    return float(bpm), filtered.astype(np.float32), float(quality)


# ----------------------------
# Session history helpers
# ----------------------------
def empty_history_df():
    return pd.DataFrame(
        columns=["timestamp", "bpm", "quality", "bp_sys", "bp_dia", "spo2", "note"]
    )


# ----------------------------
# Gradio callbacks
# ----------------------------
def process_video(
    video_file, save_history, bp_sys, bp_dia, spo2, note, state_df: pd.DataFrame
):
    t0 = time.time()
    if video_file is None:
        return (
            "No video provided. Record 8–15s with fingertip over camera + flash.",
            None,
            state_df if state_df is not None else empty_history_df(),
        )

    # Persist upload to a temp file so OpenCV can read it
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.write(video_file.read())
    tmp.flush()
    tmp.close()
    path = tmp.name

    try:
        trace, fps, duration = video_to_mean_green_trace(path)
    except Exception as e:
        return (f"Error reading video: {e}", None, state_df)

    if duration < 7:
        return ("Please record at least 7–10 seconds.", None, state_df)

    bpm, filtered, quality = estimate_bpm_and_quality(trace, fps)
    if bpm is None:
        return (
            "Could not detect a reliable heart rate. Keep finger still, steady pressure, moderate lighting.",
            None,
            state_df,
        )

    bpm_int = int(round(bpm))
    q_int = int(round(quality)) if quality is not None else None

    # Plot raw + filtered for user feedback
    fig, ax = plt.subplots(figsize=(6, 2))
    t = np.arange(len(trace)) / (fps or 30.0)
    ax.plot(t, trace, alpha=0.6, label="raw mean green")
    if filtered is not None:
        ax.plot(t, filtered, alpha=0.9, label="filtered")
    ax.set_xlabel("time (s)")
    ax.set_ylabel("signal")
    ax.set_title(f"PPG preview — {bpm_int} BPM — Quality {q_int}/100")
    ax.legend(loc="upper right")
    plt.tight_layout()

    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)

    # Update history (in-session only)
    if save_history:
        ts = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
        row = {
            "timestamp": ts,
            "bpm": bpm_int,
            "quality": q_int,
            "bp_sys": int(bp_sys) if bp_sys not in (None, "") else None,
            "bp_dia": int(bp_dia) if bp_dia not in (None, "") else None,
            "spo2": float(spo2) if spo2 not in (None, "") else None,
            "note": note or "",
        }
        state_df = pd.concat([state_df, pd.DataFrame([row])], ignore_index=True)

    elapsed = time.time() - t0
    msg = f"Estimated: {bpm_int} BPM — Quality {q_int}/100 — Processed in {elapsed:.1f}s"
    return msg, buf.getvalue(), state_df


def export_csv(state_df: pd.DataFrame):
    if state_df is None or state_df.empty:
        return None
    return ("pulsesight_history.csv", state_df.to_csv(index=False).encode("utf-8"))


# ----------------------------
# Build UI
# ----------------------------
APP_TITLE = "PulseSight"
APP_SUBTITLE = "Camera PPG heart-rate (demo) + manual vitals. Educational only."

with gr.Blocks(title=APP_TITLE, fill_height=True) as demo:
    gr.Markdown(f"# {APP_TITLE}\n**{APP_SUBTITLE}**")

    with gr.Row():
        with gr.Column(scale=2):
            video_in = gr.Video(
                source="webcam",  # allow recording from camera
                label="Record 8–15s (finger over camera + flash)",
                show_label=True,
            )
            with gr.Row():
                save_chk = gr.Checkbox(
                    label="Save this reading to session history", value=True
                )
                run_btn = gr.Button("Measure", variant="primary")
            with gr.Accordion("Manual entries (optional)", open=False):
                with gr.Row():
                    bp_sys = gr.Number(label="BP systolic (manual)", precision=0)
                    bp_dia = gr.Number(label="BP diastolic (manual)", precision=0)
                with gr.Row():
                    spo2 = gr.Number(label="SpO2 % (manual)", precision=1)
                    note = gr.Textbox(
                        label="Notes / Temp (optional)",
                        placeholder="e.g., 36.7 C or 'post exercise'",
                    )
            gr.Markdown(
                "Tips: keep finger still, moderate pressure, use flash if available, avoid heavy motion."
            )

        with gr.Column(scale=1):
            result_box = gr.Textbox(label="Result & feedback", interactive=False)
            preview_img = gr.Image(label="Signal preview (raw + filtered)")

            gr.Markdown("### Session history")
            state_df = gr.State(empty_history_df())
            history_table = gr.Dataframe(
                value=empty_history_df(), interactive=False, label="History"
            )
            with gr.Row():
                export_btn = gr.Button("Export CSV")
                export_file = gr.File(label="Download CSV")
                clear_btn = gr.Button("Clear history")

    gr.Markdown(
        "---\n"
        "**Disclaimer:** PulseSight provides approximate readings for informational and educational "
        "purposes only and is not a medical device."
    )

    # Wire callbacks
    def on_measure(
        video_file, save_flag, s, d, sp, nt, hist_df: pd.DataFrame
    ):
        msg, img_bytes, new_df = process_video(
            video_file, save_flag, s, d, sp, nt, hist_df
        )
        return msg, img_bytes, new_df, new_df

    run_btn.click(
        on_measure,
        inputs=[video_in, save_chk, bp_sys, bp_dia, spo2, note, state_df],
        outputs=[result_box, preview_img, history_table, state_df],
    )

    export_btn.click(export_csv, inputs=[state_df], outputs=[export_file])
    clear_btn.click(lambda _: empty_history_df(), inputs=[state_df], outputs=[state_df, history_table])


if __name__ == "__main__":
    # Works on Hugging Face Spaces without share=True
    demo.launch(server_name="0.0.0.0", server_port=7860)