import gradio as gr
import pandas as pd
import numpy as np
import joblib
from skimage.measure import shannon_entropy
from skimage.color import rgb2hsv
from scipy.ndimage import generic_filter
import cv2
from PIL import Image
from sklearn.preprocessing import LabelEncoder

# Load trained model and scaler
model = joblib.load('lgbm_model.pkl')  # Replace with actual path
scaler = joblib.load('minmax_scaler.pkl')  # Replace with actual path

# Define the expected feature names manually
expected_features = ['meanr', 'meang', 'meanb', 'HHR', 'Ent', 'B', 
                     'g1', 'g2', 'g3', 'g4', 'g5', 'Age']  # No 'Hgb'

# Include 'Gender' if it was used during training
use_gender = True  # Set to False if your model was not trained with 'Gender'
if use_gender:
    expected_features.append('Gender')

# Initialize LabelEncoder for gender encoding (if used in training)
gender_encoder = LabelEncoder()
gender_encoder.fit(['Female', 'Male'])

# Function to extract features
def extract_features(image):
    image = np.array(image)

    meanr = np.mean(image[:, :, 0])
    meang = np.mean(image[:, :, 1])
    meanb = np.mean(image[:, :, 2])

    hsv_image = rgb2hsv(image)
    hue = hsv_image[:, :, 0]
    high_hue_pixels = np.sum(hue > 0.95)
    total_pixels = hue.size
    HHR = high_hue_pixels / total_pixels

    gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    Ent = shannon_entropy(gray_image)
    B = np.mean(gray_image)

    def g1_filter(window): return window[4] - np.min(window)
    def g2_filter(window): return np.max(window) - window[4]
    def g3_filter(window): return window[4] - np.mean(window)
    def g4_filter(window): return np.std(window)
    def g5_filter(window): return window[4]

    g1 = generic_filter(gray_image, g1_filter, size=3).mean()
    g2 = generic_filter(gray_image, g2_filter, size=3).mean()
    g3 = generic_filter(gray_image, g3_filter, size=3).mean()
    g4 = generic_filter(gray_image, g4_filter, size=3).mean()
    g5 = generic_filter(gray_image, g5_filter, size=3).mean()

    return {
        "meanr": meanr, "meang": meang, "meanb": meanb,
        "HHR": HHR, "Ent": Ent, "B": B, "g1": g1,
        "g2": g2, "g3": g3, "g4": g4, "g5": g5,
    }

# Prediction function
def predict_hemoglobin(age, gender, image):
    try:
        if image is None:
            return "Error: No image uploaded. Please upload an image."

        if not isinstance(image, Image.Image):
            return "Error: Invalid image format. Please upload a valid image file."

        # Extract features
        features = extract_features(image)

        # Encode gender only if used in training
        if use_gender:
            features["Gender"] = gender_encoder.transform([gender])[0]

        features["Age"] = age

        # Convert to DataFrame
        features_df = pd.DataFrame([features])

        # Ensure only model-expected features are used
        for col in expected_features:
            if col not in features_df:
                features_df[col] = 0  # Add missing columns with default value

        features_df = features_df[expected_features]  # Ensure correct column order

        # Apply scaling
        features_scaled = scaler.transform(features_df)

        # Predict hemoglobin
        hemoglobin = model.predict(features_scaled)[0]

        return f"Predicted Hemoglobin Value: {hemoglobin:.2f}"

    except Exception as e:
        print(f"Error during prediction: {e}")
        return "An error occurred. Please check inputs and try again."

# Gradio interface
with gr.Blocks() as anemia_detection_app:
    gr.Markdown("# Hemoglobin Prediction App")

    with gr.Row():
        age_input = gr.Number(label="Age", value=25)
        gender_input = gr.Radio(label="Gender", choices=["Male", "Female"], value="Male", type="value")
    
    image_input = gr.Image(label="Upload Retinal Image", type="pil")
    output_text = gr.Textbox(label="Predicted Hemoglobin Value")

    predict_button = gr.Button("Predict")

    predict_button.click(
        fn=predict_hemoglobin,
        inputs=[age_input, gender_input, image_input],
        outputs=output_text
    )

# Run the app
if __name__ == "__main__":
    anemia_detection_app.launch(share=True)  # Enable public link