# app.py
import streamlit as st
from transformers import pipeline
from PIL import Image, ImageDraw
import torch

# Configuration de la page
st.set_page_config(
    page_title="Fraktur Detektion",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# CSS optimisé
st.markdown("""
<style>
    /* Réinitialisation complète */
    .stApp {
        background: transparent !important;
        padding: 0 !important;
    }
    
    .block-container {
        padding: 0.5rem !important;
        max-width: 100% !important;
    }
    
    /* Suppression des éléments superflus */
    #MainMenu, footer, header, .viewerBadge_container__1QSob {
        display: none !important;
    }
    
    .stDeployButton {
        display: none !important;
    }
    
    /* Style compact */
    .uploadedFile {
        border: 1px dashed var(--border-color);
        border-radius: 0.5rem;
        padding: 0.5rem;
    }
    
    .st-emotion-cache-1kyxreq {
        margin-top: -2rem !important;
    }
    
    /* Conteneurs de résultats */
    .result-box {
        padding: 0.5rem;
        border-radius: 0.375rem;
        margin: 0.25rem 0;
        border: 1px solid var(--border-color);
        background: var(--background-color);
    }
    
    /* Tabs plus compacts */
    .stTabs [data-baseweb="tab-list"] {
        gap: 0.5rem;
    }
    
    .stTabs [data-baseweb="tab"] {
        padding: 0.25rem 0.5rem;
        font-size: 0.875rem;
    }
    
    /* Variables CSS pour le thème */
    :root[data-theme="light"] {
        --background-color: rgba(249, 250, 251, 0.8);
        --border-color: #e5e7eb;
        --text-color: #1f2937;
    }
    
    :root[data-theme="dark"] {
        --background-color: rgba(17, 24, 39, 0.8);
        --border-color: #374151;
        --text-color: #e5e7eb;
    }
    
    /* Ajustements responsifs */
    @media (max-width: 768px) {
        .block-container {
            padding: 0.25rem !important;
        }
    }
</style>
<script>
function updateTheme(isDark) {
    document.documentElement.setAttribute('data-theme', isDark ? 'dark' : 'light');
}

window.addEventListener('message', function(e) {
    if (e.data.type === 'theme-change') {
        updateTheme(e.data.theme === 'dark');
    }
});

// Thème initial basé sur les préférences système
updateTheme(window.matchMedia('(prefers-color-scheme: dark)').matches);
</script>
""", unsafe_allow_html=True)

@st.cache_resource
def load_models():
    return {
        "D3STRON": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
        "Heem2": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
        "Nandodeomkar": pipeline("image-classification", 
            model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
    }

def translate_label(label):
    translations = {
        "fracture": "Knochenbruch",
        "no fracture": "Kein Bruch",
        "normal": "Normal",
        "abnormal": "Abnormal"
    }
    return translations.get(label.lower(), label)

def draw_boxes(image, predictions):
    draw = ImageDraw.Draw(image)
    for pred in predictions:
        box = pred['box']
        label = f"{translate_label(pred['label'])} ({pred['score']:.2%})"
        color = "#2563eb" if pred['score'] > 0.7 else "#eab308"
        
        draw.rectangle(
            [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
            outline=color,
            width=2
        )
        
        text_bbox = draw.textbbox((box['xmin'], box['ymin']-15), label)
        draw.rectangle(text_bbox, fill=color)
        draw.text((box['xmin'], box['ymin']-15), label, fill="white")
    
    return image

def main():
    models = load_models()
    
    # Contrôle de confiance simplifié
    conf_threshold = st.slider(
        "Konfidenzschwelle",
        min_value=0.0,
        max_value=1.0,
        value=0.60,
        step=0.05,
        help="Schwellenwert für die Erkennung (0-1)"
    )
    
    # Upload plus propre
    uploaded_file = st.file_uploader(
        "",
        type=['png', 'jpg', 'jpeg'],
        key="xray_upload"
    )

    if uploaded_file:
        col1, col2 = st.columns([1, 1])
        
        with col1:
            image = Image.open(uploaded_file)
            max_size = (300, 300)
            image.thumbnail(max_size, Image.Resampling.LANCZOS)
            st.image(image, use_container_width=True)

        with col2:
            tab1, tab2 = st.tabs(["📊 Klassifizierung", "🔍 Lokalisierung"])
            
            with tab1:
                for name in ["Heem2", "Nandodeomkar"]:
                    with st.spinner("Analyse..."):
                        predictions = models[name](image)
                        for pred in predictions:
                            if pred['score'] >= conf_threshold:
                                score_color = "#22c55e" if pred['score'] > 0.7 else "#eab308"
                                st.markdown(f"""
                                    <div class='result-box'>
                                        <span style='color: {score_color}; font-weight: 500;'>
                                            {pred['score']:.1%}
                                        </span> - {translate_label(pred['label'])}
                                    </div>
                                """, unsafe_allow_html=True)

            with tab2:
                with st.spinner("Lokalisierung..."):
                    predictions = models["D3STRON"](image)
                    filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
                    
                    if filtered_preds:
                        result_image = image.copy()
                        result_image = draw_boxes(result_image, filtered_preds)
                        st.image(result_image, use_container_width=True)
                        
                        for pred in filtered_preds:
                            st.markdown(f"""
                                <div class='result-box'>
                                    {translate_label(pred['label'])}: {pred['score']:.1%}
                                </div>
                            """, unsafe_allow_html=True)
                    else:
                        st.info("Keine Erkennungen über dem Schwellenwert")
    else:
        st.info("Röntgenbild hochladen (JPEG, PNG, max. 5MB)")

if __name__ == "__main__":
    main()