import gradio as gr
import numpy as np
import tensorflow as tf
from scipy.signal import find_peaks

model = tf.keras.models.load_model('aritmi_distilled.h5')

def extract_features(ekg_signal):
    if len(ekg_signal) < 10:
        return None

    r_peaks, _ = find_peaks(ekg_signal, height=0.5, distance=60)

    if len(r_peaks) < 2:
        return None

    features = []
    rr_intervals = np.diff(r_peaks)

    if len(rr_intervals) == 0:
        return None

    features.append(np.nanmean(rr_intervals))
    features.append(np.nanstd(rr_intervals))
    features.append(np.nanmean(ekg_signal))
    features.append(np.nanmean(r_peaks))
    features.append(np.nanmean(np.diff(r_peaks)))

    pq_interval = np.nanmean(rr_intervals)
    qt_interval = np.nanmean(rr_intervals)
    st_interval = np.nanmean(rr_intervals)

    features.extend([pq_interval, qt_interval, st_interval])
    return features

def predict(ekg_signal):
    try:
        ekg_signal_list = [float(value.strip()) for value in ekg_signal.split(',') if value.strip()]
        ekg_signal_array = np.array(ekg_signal_list, dtype=np.float32)
    except Exception as e:
        return f"Invalid dtype. Error: {str(e)}"

    if ekg_signal_array.size == 0:
        return "Empty ECG signal."

    features = extract_features(ekg_signal_array)
    if features is None:
        return "Couldn't extract feature."

    features = np.array(features).reshape(1, -1)
    prediction = model.predict(features)
    risk_probabilities = prediction[0]
    aritmi_riski_yuzde = risk_probabilities[1] * 100

    return f"%{aritmi_riski_yuzde:.2f}"



interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="ECG Signal",
        placeholder="0.034, 0.181, 0.264, ...",
    ),
    outputs="text",
    description="Arrhythmia risk predicting model from ECG data",
    live=False
)



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