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

def extract_ppg_signal(video_file):
    cap = cv2.VideoCapture(video_file)
    ppg_signal = []
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        ppg_signal.append(np.mean(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)))
    cap.release()
    return np.array(ppg_signal)

def process_ppg_signal(ppg_signal, fs=30):
    nyquist = 0.5 * fs
    b, a = butter(1, [0.5 / nyquist, 3.0 / nyquist], btype='band')
    return filtfilt(b, a, ppg_signal)

def detect_hrv(video_file):
    filtered_signal = process_ppg_signal(extract_ppg_signal(video_file))
    peaks, _ = find_peaks(filtered_signal, distance=15)  # 30 fps / 2.5
    rr_intervals = np.diff(peaks) / 30
    heart_rate = 60 / rr_intervals.mean()
    hrv = np.std(rr_intervals)
    
    plt.figure(figsize=(10, 4))
    time = np.arange(len(filtered_signal)) / 30
    plt.plot(time, filtered_signal, label='Filtered PPG Signal')
    plt.plot(time[peaks], filtered_signal[peaks], 'ro', label='Detected Peaks')
    plt.title('Heart Rate Variability over Time')
    plt.xlabel('Time (s)')
    plt.ylabel('PPG Signal Intensity')
    plt.legend()
    plt.tight_layout()
    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
        plt.savefig(temp_file.name)
    plt.close()
    
    return f"{hrv:.2f} ms", f"{heart_rate:.2f} BPM", temp_file.name

def create_hrv_tab():
    with gr.Row():
        with gr.Column(scale=2):
            input_video = gr.Video(label="Input Video")
            with gr.Row():
                clear_btn = gr.Button("Clear", scale=1)
                submit_btn = gr.Button("Analyze", scale=1, elem_classes="submit")
        with gr.Column(scale=1):
            output_hrv = gr.Label(label="HRV Value")
            output_hr = gr.Label(label="Average Heart Rate")
            output_plot = gr.Image(label="HRV Plot")
    
    submit_btn.click(detect_hrv, inputs=[input_video], outputs=[output_hrv, output_hr, output_plot], queue=True)
    clear_btn.click(lambda: (None, None, None, None), outputs=[input_video, output_hrv, output_hr, output_plot], queue=True)
    gr.Examples(["./assets/videos/fitness.mp4"], [input_video])