import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image, ImageStat
import cv2

# 🧠 Load your TFLite model
interpreter = tf.lite.Interpreter(model_path="stool_model.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 🏷️ Labels of your stool model
labels = ["bloody", "hard stool", "normal", "parasite", "watery"]

# 💬 Diagnosis summaries
diagnosis_advice = {
    "bloody": "Possible hemorrhagic gastroenteritis. ⚠️ Visit a vet immediately as it may indicate internal bleeding or infection.",
    "hard stool": "Your pet may be constipated. 💧 Encourage hydration and increase dietary fiber intake.",
    "normal": "✅ Healthy stool detected. Your pet appears to be in good digestive health.",
    "parasite": "⚠️ Possible tapeworm or other parasite infection. Schedule deworming and a vet check-up.",
    "watery": "Possible diarrhea. 💧 Ensure hydration and monitor for dehydration or weakness.",
    "Not stool image": "❌ This image doesn’t appear to be a stool sample. Please upload a clearer photo of stool.",
    "Uncertain / unclear stool image": "⚠️ The image is unclear or confidence is low. Try retaking the photo in better lighting."
}

# 🧩 Image preprocessing for the TFLite model
def preprocess_image(img: Image.Image):
    img = img.convert("RGB").resize((128, 128))
    arr = np.asarray(img).astype(np.float32) / 255.0
    arr = np.expand_dims(arr, axis=0)
    return arr

# 🚫 Strict non-stool detection
def is_not_stool_image(image):
    """Multi-layer detection for non-stool photos"""
    np_img = np.array(image.convert("RGB"))
    stat = ImageStat.Stat(image)

    # 🧠 1️⃣ Brightness & contrast checks
    brightness = stat.mean[0]
    contrast = stat.stddev[0]
    if brightness > 210 or brightness < 30:
        return True  # too bright or dark
    if contrast < 25:
        return True  # too flat / no texture

    # 🧠 2️⃣ Dominant color detection (stool tends to be brownish)
    avg_color = np.mean(np_img.reshape(-1, 3), axis=0)
    r, g, b = avg_color

    # Detect unnatural colors (too blue, green, or red)
    if (r > 200 and g < 100 and b < 100) or (g > 180 and r < 100) or (b > 180):
        return True

    # 🧠 3️⃣ Color tone validation — real stool is usually brown (RGB ratios)
    brownish_score = (r * 0.5 + g * 0.3 + b * 0.2)
    grayish_score = abs(r - g) + abs(g - b)
    if brownish_score < 60 or grayish_score > 80:
        return True

    # 🧠 4️⃣ Texture detection with OpenCV
    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
    lap_var = cv2.Laplacian(gray, cv2.CV_64F).var()  # blur measure
    if lap_var < 50:
        return True  # image too smooth or blurry (likely not stool)

    return False  # passes all checks, likely stool

# 🧠 Main classifier
def classify_image(image):
    try:
        # 🧩 Step 1: Reject non-stool images first
        if is_not_stool_image(image):
            return {"Not stool image": 1.0}, diagnosis_advice["Not stool image"]

        # 🧩 Step 2: Run TFLite model
        input_data = preprocess_image(image)
        interpreter.set_tensor(input_details[0]['index'], input_data)
        interpreter.invoke()
        output_data = interpreter.get_tensor(output_details[0]['index'])[0]

        # 🧩 Step 3: Sort predictions
        results = {labels[i]: float(output_data[i]) for i in range(len(labels))}
        sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
        top_label, top_score = list(sorted_results.items())[0]

        # 🧩 Step 4: Handle low-confidence predictions
        if top_score < 0.45:
            return {"Uncertain / unclear stool image": top_score}, diagnosis_advice["Uncertain / unclear stool image"]

        # 🧩 Step 5: Add readable confidence %
        formatted_results = {
            f"{label} ({score * 100:.2f}%)": score for label, score in sorted_results.items()
        }

        # 🩺 Step 6: Generate human-readable advice
        advice = diagnosis_advice.get(top_label, "No advice available for this diagnosis.")
        return formatted_results, advice

    except Exception as e:
        return {"Error": 0.0}, f"⚠️ Error: {str(e)}"

# 🎨 Gradio Interface
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="📸 Upload Stool Image"),
    outputs=[
        gr.Label(num_top_classes=3, label="Predicted Diagnosis & Confidence"),
        gr.Textbox(label="💬 Diagnosis Summary", lines=3),
    ],
    title="🐾 Pet Stool Diagnosis AI (Strict Validation)",
    description=(
        "Upload a stool image for AI-based diagnosis. The system will first verify if it's a real stool photo, "
        "then predict the stool type and give a simple health summary."
    ),
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)