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

IMG_HEIGHT = 128
IMG_WIDTH = 128

st.set_page_config(
    page_title="BSDS500 Denoising Autoencoder",
    layout="centered"
)

st.title("🧼 BSDS500 Denoising Autoencoder")
st.write(
    """
    Deze app gebruikt een **convolutionele autoencoder** om ruis uit afbeeldingen te halen.
    Upload een foto (JPG/PNG), kies hoeveel ruis je wilt toevoegen, en bekijk het resultaat.
    """
)

@st.cache_resource(show_spinner=True)
def load_model():
    # model staat in dezelfde map als dit bestand (src/)
    current_dir = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(current_dir, "bsds500_denoising_autoencoder.h5")

    if not os.path.exists(model_path):
        st.error(f"Modelbestand niet gevonden op pad:\n`{model_path}`")
        return None

    try:
        model = tf.keras.models.load_model(model_path, compile=False)
        return model
    except Exception as e:
        st.error(f"❌ Kon het model niet laden: {e}")
        return None

model = load_model()
if model is None:
    st.stop()

def preprocess_image(image: Image.Image):
    image = image.convert("RGB")
    image = image.resize((IMG_WIDTH, IMG_HEIGHT))
    img_array = np.array(image).astype("float32") / 255.0
    return img_array

def add_noise(img_array: np.ndarray, noise_factor: float):
    noise = np.random.normal(0.0, 1.0, img_array.shape)
    noisy = img_array + noise_factor * noise
    return np.clip(noisy, 0.0, 1.0)

def denoise_image(noisy_array: np.ndarray):
    inp = np.expand_dims(noisy_array, axis=0)
    reconstructed = model.predict(inp, verbose=0)[0]
    return reconstructed

uploaded_file = st.file_uploader(
    "📁 Upload een afbeelding",
    type=["jpg", "jpeg", "png"]
)

noise_factor = st.slider(
    "📉 Hoeveel ruis toevoegen?",
    min_value=0.0,
    max_value=0.5,
    value=0.1,
    step=0.01
)

run_button = st.button("✨ Maak reconstructie")

if uploaded_file is not None:
    pil_image = Image.open(uploaded_file)
    clean_img = preprocess_image(pil_image)

    col1, col2, col3 = st.columns(3)

    with col1:
        st.subheader("Origineel")
        st.image(clean_img, use_container_width=True)

    if run_button:
        noisy_img = add_noise(clean_img, noise_factor)
        reconstructed_img = denoise_image(noisy_img)

        with col2:
            st.subheader("Noisy")
            st.image(noisy_img, use_container_width=True)

        with col3:
            st.subheader("Reconstructed")
            st.image(reconstructed_img, use_container_width=True)
    else:
        st.info("Klik op **'✨ Maak reconstructie'** om de autoencoder te gebruiken.")
else:
    st.warning("Upload eerst een afbeelding (JPG/PNG).")