import gradio as gr
import numpy as np
import pandas as pd
from io import BytesIO
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.signal import find_peaks, resample
from PIL import Image

# ----------- LOAD MODEL -----------
model = tf.keras.models.load_model("ecg_cnn_model.h5", compile=False)

# Auto-detect what input length the model expects
MODEL_INPUT_SHAPE = model.input_shape  # e.g. (None, 187) or (None, 6016) etc.
if isinstance(MODEL_INPUT_SHAPE, list):
    MODEL_INPUT_SHAPE = MODEL_INPUT_SHAPE[0]
EXPECTED_LEN = MODEL_INPUT_SHAPE[-1]
print(f"[CardioSense] Model expects input length: {EXPECTED_LEN}")

# ----------- PREPROCESS -----------
def preprocess(filepath):
    try:
        data = pd.read_csv(filepath, encoding='latin1', on_bad_lines='skip')
    except Exception:
        return None

    data.columns = data.columns.str.replace("'", "").str.strip().str.upper()

    if data.shape[1] < 1:
        return None

    if 'MLII' in data.columns:
        signal = data['MLII'].dropna().values
    elif data.shape[1] >= 2:
        signal = data.iloc[:, 1].dropna().values
    else:
        signal = data.iloc[:, 0].dropna().values

    if len(signal) < 100:
        return None

    signal = signal.astype(np.float32)
    signal = (signal - np.mean(signal)) / (np.std(signal) + 1e-8)

    # Resample raw signal to exactly what the model expects
    raw_segment = signal[:min(len(signal), EXPECTED_LEN * 4)]  # grab a generous chunk
    resampled = resample(raw_segment, EXPECTED_LEN).astype(np.float32)
    segment = resampled.reshape(1, EXPECTED_LEN)

    return segment, signal

# ----------- PREDICT -----------
def predict_ecg(file):
    if file is None:
        return "❌ No file uploaded.", None

    result_data = preprocess(file.name)
    if result_data is None:
        return "❌ Invalid ECG file. Make sure it has at least 100 rows of numeric data.", None

    segment, signal = result_data

    try:
        pred = model.predict(segment, verbose=0)[0][0]
    except Exception as e:
        return f"❌ Prediction failed: {str(e)}", None

    confidence = round(float(pred) * 100, 2) if pred > 0.5 else round((1 - float(pred)) * 100, 2)
    result = "Abnormal" if pred > 0.5 else "Normal"

    # Heart rate estimation
    peaks, _ = find_peaks(signal, distance=150, height=0.5)
    if len(peaks) > 1 and np.mean(np.diff(peaks)) > 0:
        rr = np.diff(peaks) / 360.0
        heart_rate = int(60 / np.mean(rr))
        heart_rate = max(30, min(220, heart_rate))  # clamp to physiological range
    else:
        heart_rate = None

    # Plot
    plot_len = min(500, len(signal))
    plt.figure(figsize=(8, 3))
    plt.plot(signal[:plot_len], color='royalblue', label="ECG Signal", linewidth=0.8)
    for p in peaks:
        if p < plot_len:
            plt.plot(p, signal[p], "ro", markersize=4)
    bpm_label = f"{heart_rate} BPM" if heart_rate else "-- BPM"
    plt.title(f"ECG Signal ({bpm_label})", fontsize=13)
    plt.xlabel("Sample")
    plt.ylabel("Amplitude")
    plt.grid(alpha=0.3)
    plt.tight_layout()

    img_buf = BytesIO()
    plt.savefig(img_buf, format='png', dpi=120)
    img_buf.seek(0)
    plt.close()
    pil_img = Image.open(img_buf)

    emoji = "🔴" if result == "Abnormal" else "🟢"
    advice = "Please consult a cardiologist." if result == "Abnormal" else "Your ECG looks healthy."
    output_text = (
        f"{emoji} **Result: {result}**\n\n"
        f"📊 Confidence: {confidence}%\n\n"
        f"❤️ Heart Rate: {heart_rate if heart_rate else 'N/A'} BPM\n\n"
        f"⚠️ {advice}\n\n"
        f"_This tool is for educational purposes only. Always consult a doctor._"
    )
    return output_text, pil_img

# ----------- GRADIO UI -----------
with gr.Blocks(title="CardioSense AI") as demo:
    gr.Markdown("""
    # 💓 CardioSense AI
    ### ECG Heart Condition Detector
    Upload your ECG data as a **CSV file** to detect if it is **Normal** or **Abnormal**.
    > Supports MIT-BIH format CSVs (e.g. with MLII column) or any single-column numeric ECG data.
    """)

    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="📂 Upload ECG CSV File", file_types=[".csv"])
            predict_btn = gr.Button("🔍 Analyze ECG", variant="primary")
        with gr.Column():
            result_output = gr.Markdown(label="Result")
            plot_output = gr.Image(label="📈 ECG Graph", type="pil")

    predict_btn.click(
        fn=predict_ecg,
        inputs=file_input,
        outputs=[result_output, plot_output]
    )

    gr.Markdown("> ⚠️ For educational use only. Not a substitute for medical advice.")

demo.launch()