import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
import joblib

# Load the trained model and scaler
model = load_model("seizure_model.h5")
scaler = joblib.load("scaler.pkl")

def process_input(raw_input):
    # Handle pasted input (space, tab, newline separated)
    try:
        eeg_array = np.array([float(x) for x in raw_input.strip().split()])
    except:
        return "❌ Invalid input: Please enter only numbers separated by spaces, tabs, or newlines."

    if len(eeg_array) > 178:
        eeg_array = eeg_array[:178]
    elif len(eeg_array) < 178:
        return "❌ Input must contain at least 178 EEG values."

    eeg_scaled = scaler.transform(eeg_array.reshape(1, -1))
    eeg_cnn_input = eeg_scaled.reshape(1, 178, 1)

    prediction = model.predict(eeg_cnn_input)[0].argmax() + 1

    if prediction == 1:
        return "🚨 Seizure Detected (Class 1)"
    else:
        return f"✅ Normal EEG (Class {prediction})"

def predict_from_file(file):
    raw = file.read().decode("utf-8")
    return process_input(raw)

def predict_from_text(raw_input):
    return process_input(raw_input)

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft(), title="RakshAI - Seizure Detection") as app:
    gr.Markdown("## 🧠 RakshAI - Seizure Detection")
    gr.Markdown("Upload a `.txt` file or paste **178 EEG values** separated by spaces, tabs, or newlines.")

    with gr.Row():
        textbox = gr.Textbox(lines=10, label="📝 Paste EEG Values")
        file = gr.File(label="📂 Upload EEG File (.txt)", file_types=[".txt"])

    predict_btn = gr.Button("🚀 Predict Seizure Status")

    output = gr.Textbox(label="🔮 Prediction Output")

    predict_btn.click(fn=predict_from_text, inputs=textbox, outputs=output)
    file.change(fn=predict_from_file, inputs=file, outputs=output)

app.launch()