import gradio as gr
import numpy as np

# Load trained model parameters (assumed to be saved as .npy file)
theta_final = np.load("theta_final.npy")  # Save your trained theta and upload it to HF Spaces

# Mean and standard deviation for normalization (precomputed from dataset)
mu = np.array([0.5] * 16)  # Replace with actual values
sigma = np.array([0.2] * 16)  # Replace with actual values
sigma = np.where(sigma == 0, 1, sigma)  # Prevent division by zero

def predict_stroke_risk(*symptoms):
    user_input = np.array([symptoms])
    normalized_input = (user_input - mu) / sigma
    normalized_input = np.c_[np.ones((normalized_input.shape[0], 1)), normalized_input]  # Add bias term
    y_pred = normalized_input @ theta_final
    y_pred = np.clip(y_pred[0][0], 0, 100)  # Ensure output is between 0% and 100%
    return f"Predicted Stroke Risk: {y_pred:.2f}%"

# Symptoms list
symptoms = [
    "Chest Pain", "Shortness of Breath", "Irregular Heartbeat", "Fatigue/Weakness", 
    "Dizziness", "Swelling/Edema", "Pain in Neck/Jaw/Shoulder/Back", "Excessive Sweating", 
    "Persistent Cough", "Nausea/Vomiting", "High Blood Pressure", "Chest Discomfort During Activity", 
    "Cold Hands/Feet", "Snoring/Sleep Apnea", "Anxiety/Feeling of Doom", "Age > 50"
]

with gr.Blocks(css="body {font-family: monospace; text-align: center; color: #3A78F2;}") as demo:
    gr.Markdown("# Stroke Prediction")
    gr.Markdown("### Select your symptoms:")
    
    checkbox_list = []
    with gr.Row():
        for i in range(0, len(symptoms), 5):  # 5 symptoms per column
            with gr.Column():
                for j in range(5):
                    if i + j < len(symptoms):
                        checkbox_list.append(gr.Checkbox(label=symptoms[i + j], value=False))

    submit_button = gr.Button("Submit")
    output_text = gr.Markdown("### Possibility of stroke: ______")

    submit_button.click(fn=predict_stroke_risk, inputs=checkbox_list, outputs=output_text)

demo.launch()