import os
# --- LIGNE DE SAUVETAGE OBLIGATOIRE ---
# Cela corrige l'erreur "Descriptors cannot be created directly"
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
from PIL import Image

# ==========================================
# 1. CONFIGURATION
# ==========================================
st.set_page_config(page_title="Détection Malaria", layout="wide")
IMG_SIZE = (64, 64)

# ==========================================
# 2. CHARGEMENT DU MODÈLE
# ==========================================
@st.cache_resource
def load_model():
    try:
        # TF 2.12 lit nativement 'batch_shape' sans planter
        model = tf.keras.models.load_model('malaria_model_finetuned.h5', compile=False)
        return model
    except Exception as e:
        st.error(f"Erreur fatale : {e}")
        return None

model = load_model()

# ==========================================
# 3. PRÉTRAITEMENT
# ==========================================
def biological_preprocessing_inference(image_pil):
    img_np = np.array(image_pil)
    img_res = cv2.resize(img_np, IMG_SIZE)

    # 7 Canaux
    img_norm = img_res / 255.0
    gray = cv2.cvtColor(img_res, cv2.COLOR_RGB2GRAY)
    _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    gray_no_bg = cv2.bitwise_and(gray, gray, mask=mask) / 255.0
    canny = cv2.Canny(gray, 40, 120) / 255.0
    
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    sobel_norm = cv2.magnitude(sobelx, sobely)
    sobel_norm = cv2.normalize(sobel_norm, None, 0, 1, cv2.NORM_MINMAX)

    img_7ch = np.dstack((img_norm, gray/255.0, gray_no_bg, canny, sobel_norm)).astype(np.float32)

    # Bio Features
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    area, perim, circ, rect_area = 0.0, 0.0, 0.0, 0.0
    if contours:
        cnt = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(cnt) / (IMG_SIZE[0] * IMG_SIZE[1])
        perim = cv2.arcLength(cnt, True) / IMG_SIZE[0]
        if perim > 0: circ = (4 * np.pi * area) / (perim**2)
        x, y, w, h = cv2.boundingRect(cnt)
        if w*h > 0: rect_area = area / (w * h / (IMG_SIZE[0] * IMG_SIZE[1]))

    bio_desc = np.array([area, perim, circ, rect_area], dtype=np.float32)

    return np.expand_dims(img_7ch, axis=0), np.expand_dims(bio_desc, axis=0)

# ==========================================
# 4. INTERFACE
# ==========================================
st.title("🔬 Détection Malaria")
uploaded_file = st.file_uploader("Image", type=["png", "jpg", "jpeg"])

if uploaded_file and model:
    image = Image.open(uploaded_file).convert('RGB')
    col1, col2 = st.columns(2)
    with col1:
        st.image(image, caption="Image source", width=300)
    
    if st.button("Lancer le diagnostic"):
        with st.spinner("Analyse..."):
            try:
                img_in, bio_in = biological_preprocessing_inference(image)
                # Double sécurité pour la prédiction
                try:
                    pred = model.predict({'img_input': img_in, 'bio_input': bio_in})
                except:
                    pred = model.predict([img_in, bio_in])
                
                idx = np.argmax(pred[0])
                conf = np.max(pred[0])
                label = "Infecté (Parasitized) 🦠" if idx == 0 else "Sain (Uninfected) 🛡️"
                
                with col2:
                    if idx == 0: st.error(f"### {label}")
                    else: st.success(f"### {label}")
                    st.metric("Confiance", f"{conf:.2%}")
            except Exception as e:
                st.error(f"Erreur : {e}")