import gradio as gr
import cv2
import numpy as np
from PIL import Image


IMG_SIZE = 256
REFERENCE_PATH = "reference_bottle.npy"

THRESHOLD_SCORE = 400

reference = np.load(REFERENCE_PATH)


def detect_classical(input_image):
    if input_image is None:
        return "<h2>Veuillez importer une image.</h2>", None, None, None

    img_rgb = np.array(input_image.convert("RGB"))
    img_rgb = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))

    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)

    diff = cv2.absdiff(gray, reference)

    diff_blur = cv2.GaussianBlur(diff, (5, 5), 0)

    _, mask = cv2.threshold(diff_blur, 35, 255, cv2.THRESH_BINARY)

    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    anomaly_score = int(np.sum(mask > 0))

    if anomaly_score > THRESHOLD_SCORE:
        prediction = "DÉFECTUEUSE"
        color = "#c0392b"
        emoji = "❌"
    else:
        prediction = "BONNE"
        color = "#27ae60"
        emoji = "✅"

    result_html = f"""
    <div style="
        text-align:center;
        padding:25px;
        border-radius:18px;
        background-color:#f8f9fa;
        border:3px solid {color};
    ">
        <h1 style="color:{color}; font-size:44px;">
            {emoji} Prédiction : {prediction}
        </h1>

        <p style="font-size:24px;">
            Score d'anomalie : <b>{anomaly_score}</b>
        </p>

        <p style="font-size:20px;">
            Seuil utilisé : <b>{THRESHOLD_SCORE}</b>
        </p>
    </div>
    """

    return result_html, img_rgb, diff, mask


custom_css = """
.gradio-container {
    max-width: 1250px !important;
    margin: auto !important;
}

h1 {
    font-size: 26px !important;
    margin-bottom: 5px !important;
}

.compact-text {
    font-size: 14px !important;
    margin-bottom: 5px !important;
}

footer {
    display: none !important;
}
"""

example_paths = [
    "examples/good_141.png",
    "examples/good_093.png",
    "examples/good_146.png",
    "examples/defect_017bd.png",
    "examples/defect_004sd.png",
    "examples/defect_010c.png",
    "examples/defect_016sd.png",
]


def load_gallery_example(evt: gr.SelectData):
    image_path = example_paths[evt.index]
    image = Image.open(image_path).convert("RGB")
    result_html, original, diff, mask = detect_classical(image)
    return image, result_html, original, diff, mask


with gr.Blocks(
    title="Détection de défauts - Méthode classique",
    css=custom_css
) as demo:

    gr.Markdown(
        """
        # Détection de défauts sur bouteilles - Méthode classique
        <p class="compact-text">
        Comparaison avec une image de référence + score d'anomalie. Importer une image ou cliquer sur un exemple.
        </p>
        """
    )

    with gr.Row():

        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil",
                label="Image à analyser",
                height=260
            )

            button = gr.Button("Analyser l'image", variant="primary")

        with gr.Column(scale=1):
            result_output = gr.HTML(label="Résultat")

    with gr.Row():
        original_output = gr.Image(label="Image analysée", height=160)
        diff_output = gr.Image(label="Différence", height=160)
        mask_output = gr.Image(label="Masque", height=160)

    gr.Markdown("### Exemples de test")

    gallery = gr.Gallery(
        value=example_paths,
        label="Cliquer sur une image pour la tester",
        columns=7,
        height=140,
        object_fit="contain",
        allow_preview=False
    )

    button.click(
        fn=detect_classical,
        inputs=input_image,
        outputs=[
            result_output,
            original_output,
            diff_output,
            mask_output
        ]
    )

    gallery.select(
        fn=load_gallery_example,
        inputs=None,
        outputs=[
            input_image,
            result_output,
            original_output,
            diff_output,
            mask_output
        ]
    )

demo.launch()