import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

import numpy as np
import gradio as gr
import joblib
import spaces  # Still needed for the dummy function
from tensorflow.keras.models import load_model

# ============================================================
# LOAD MODELS & SCALERS
# ============================================================
knn_model = joblib.load("knn_model.pkl")
svm_model = joblib.load("rbf_svm_model.pkl")
ann_model = load_model("ann_model.keras")

scaler_rfe = joblib.load("scaler_rfe.pkl")
selected_features = joblib.load("selected_features.pkl")

# ============================================================
# ENCODING MAPS
# ============================================================
gender_map = {"Female": 0, "Male": 1}
binary_map = {"No": 0, "Yes": 1}
ordinal_map = {"Low": 0, "Medium": 1, "High": 2}
alcohol_map = {"None": 0, "Low": 1, "Medium": 2, "High": 3}

binary_features = {'smoking', 'family_hd', 'diabetes', 'high_bp', 'low_hdl', 'high_ldl'}
ordinal_features = {'exercise', 'stress', 'sugar'}

def encode_input_value(feature, value):
    if feature == 'gender': return gender_map[value]
    if feature in binary_features: return binary_map[value]
    if feature in ordinal_features: return ordinal_map[value]
    if feature == 'alcohol': return alcohol_map[value]
    return float(value)

# ============================================================
# ZERO-GPU WORKAROUND
# ============================================================
# This satisfies Hugging Face's ZeroGPU requirement without 
# crashing your actual Scikit-Learn/TensorFlow models.
@spaces.GPU
def dummy_startup():
    pass

# ============================================================
# PREDICTION FUNCTION (NO GPU DECORATOR)
# ============================================================
def predict_heart_disease(model_choice, *feature_values):
    try:
        encoded_inputs = [encode_input_value(f, v) for f, v in zip(selected_features, feature_values)]
        X_input = np.array([encoded_inputs], dtype=float)
        X_scaled = scaler_rfe.transform(X_input)

        if model_choice == "KNN":
            y_pred = knn_model.predict(X_scaled)[0]
            y_prob = knn_model.predict_proba(X_scaled)[0][1]
        elif model_choice == "SVM":
            y_pred = svm_model.predict(X_scaled)[0]
            if hasattr(svm_model, "predict_proba"):
                y_prob = svm_model.predict_proba(X_scaled)[0][1]
            else:
                score = svm_model.decision_function(X_scaled)[0]
                y_prob = 1 / (1 + np.exp(-score))
        else: # ANN
            y_prob = float(ann_model.predict(X_scaled, verbose=0)[0][0])
            y_pred = int(y_prob >= 0.5)

        status = "Yes" if int(y_pred) == 1 else "No"
        probability_text = f"{y_prob:.4f} ({y_prob * 100:.2f}%)"

        return status, probability_text, model_choice
        
    except Exception as e:
        # If an error happens, this will print it directly to the Gradio UI
        return f"Error: {str(e)}", "Error", model_choice

# ============================================================
# UI DEFINITION
# ============================================================
feature_labels = {
    'age': 'Age', 'gender': 'Gender', 'blood_pressure': 'Blood Pressure',
    'cholesterol': 'Cholesterol Level', 'exercise': 'Exercise Habits',
    'smoking': 'Smoking', 'family_hd': 'Family Heart Disease', 'diabetes': 'Diabetes',
    'bmi': 'BMI', 'high_bp': 'High Blood Pressure', 'low_hdl': 'Low HDL Cholesterol',
    'high_ldl': 'High LDL Cholesterol', 'alcohol': 'Alcohol Consumption',
    'stress': 'Stress Level', 'sleep_hours': 'Sleep Hours', 'sugar': 'Sugar Consumption',
    'triglyceride': 'Triglyceride Level', 'fasting_bs': 'Fasting Blood Sugar',
    'crp': 'CRP Level', 'homocysteine': 'Homocysteine Level'
}

def create_feature_input(feature):
    label = feature_labels.get(feature, feature)
    if feature == 'gender': return gr.Dropdown(["Female", "Male"], label=label, value="Female")
    if feature in binary_features: return gr.Dropdown(["No", "Yes"], label=label, value="No")
    if feature in ordinal_features: return gr.Dropdown(["Low", "Medium", "High"], label=label, value="Medium")
    if feature == 'alcohol': return gr.Dropdown(["None", "Low", "Medium", "High"], label=label, value="None")
    return gr.Number(label=label, value=0.0)

with gr.Blocks() as demo:
    gr.Markdown("<h2 style='text-align:center'>💓 Heart Disease Prediction 💓</h2>")
    
    with gr.Row():
        model_input = gr.Dropdown(["ANN", "SVM", "KNN"], label="Choose Model", value="ANN")

    feature_inputs = []
    for start in range(0, len(selected_features), 5):
        with gr.Row():
            for feature in selected_features[start:start + 5]:
                feature_inputs.append(create_feature_input(feature))

    predict_button = gr.Button("Predict 💡", variant="primary")

    with gr.Row():
        status_output = gr.Textbox(label="Heart Disease Status")
        prob_output = gr.Textbox(label="Prediction Probability")
        model_output = gr.Textbox(label="Model Selected")

    predict_button.click(
        fn=predict_heart_disease,
        inputs=[model_input] + feature_inputs,
        outputs=[status_output, prob_output, model_output]
    )

demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)