import os

os.environ["HOME"] = "/tmp"
os.environ["XDG_CONFIG_HOME"] = "/tmp/.config"
os.environ["XDG_CACHE_HOME"] = "/tmp/.cache"
os.environ["MPLCONFIGDIR"] = "/tmp/.config/matplotlib"
for d in ["/tmp", "/tmp/.config", "/tmp/.cache", "/tmp/.config/matplotlib"]:
    os.makedirs(d, exist_ok=True)

import streamlit as st
import tensorflow as tf
import numpy as np
from keras.preprocessing.image import img_to_array
from PIL import Image

IMG_SIZE = (224, 224)
LABELS = ["Fracture Detected", "No Fracture"]
MODEL_PATH = os.path.join(os.path.dirname(__file__), "model_weights.h5")

@st.cache_resource
def load_model():
    if not os.path.exists(MODEL_PATH):
        st.error(f"❌ Model weights not found at: `{MODEL_PATH}`")
        st.stop()

    base = tf.keras.applications.MobileNetV2(include_top=False, input_shape=(224, 224, 3), pooling="avg")
    x = tf.keras.layers.Dense(128, activation="relu")(base.output)
    out = tf.keras.layers.Dense(2, activation="softmax")(x)
    model = tf.keras.Model(inputs=base.input, outputs=out)
    model.load_weights(MODEL_PATH)
    return model

# --- Streamlit UI ---
st.title("X-ray Fracture Detection")
file = st.file_uploader("Upload an X-ray image", type=["png", "jpg", "jpeg"])

if file:
    st.image(file, use_column_width=True, caption="🖼 Uploaded X-ray")
    st.write(f"File: `{file.name}`")
    try:
        image = Image.open(file).convert("RGB").resize(IMG_SIZE)
        array = img_to_array(image) / 255.0
        array = np.expand_dims(array, axis=0)
        st.write(f"Image shape: `{array.shape}`")  # 🧪 Debug info

        model = load_model()
        pred = model.predict(array)
        st.write(f"Raw model prediction: `{pred}`")  # 🧪 Debug info

        class_idx = np.argmax(pred)
        confidence = pred[0][class_idx] * 100

        st.success(f"Prediction: **{LABELS[class_idx]}** ({confidence:.2f}%)")

    except Exception as e:
        st.error(f"Prediction failed: {e}")

else:
    st.info("Upload an image to get a prediction.")