import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from safetensors import safe_open

# ✅ Constants
IMG_SIZE = 224
CLASS_NAMES = ["Fractured", "Non-Fractured"]
SAFETENSOR_PATH = "osteologic.safetensors"

# ✅ Step 1: Rebuild architecture
def build_model():
    inputs = tf.keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    base_model = tf.keras.applications.MobileNetV2(weights=None, include_top=False, input_tensor=inputs)
    x = base_model.output
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(128, activation="relu", kernel_regularizer=tf.keras.regularizers.l2(0.001))(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    outputs = tf.keras.layers.Dense(len(CLASS_NAMES), activation="softmax")(x)
    model = tf.keras.Model(inputs, outputs)
    return model

# ✅ Step 2: Load weights from .safetensors
def load_weights(model, path=SAFETENSOR_PATH):
    with safe_open(path, framework="pt", device="cpu") as f:
        for layer in model.layers:
            if isinstance(layer, (tf.keras.layers.Conv2D, tf.keras.layers.Dense)):
                w_key = f"{layer.name}.weight"
                b_key = f"{layer.name}.bias"
                if w_key in f.keys() and b_key in f.keys():
                    weights = f.get_tensor(w_key)
                    bias = f.get_tensor(b_key)
                    # Adjust shape if needed (PyTorch → TF)
                    if isinstance(layer, tf.keras.layers.Conv2D):
                        weights = weights.transpose(2, 3, 1, 0)  # [out, in, h, w] → [h, w, in, out]
                    layer.set_weights([weights, bias])
    return model

# ✅ Step 3: Build and load model
model = build_model()
model = load_weights(model)

# ✅ Step 4: Prediction function
def predict(image: Image.Image):
    image = image.resize((IMG_SIZE, IMG_SIZE)).convert("RGB")
    arr = np.array(image) / 255.0
    arr = arr.reshape(1, IMG_SIZE, IMG_SIZE, 3)
    preds = model.predict(arr)[0]
    label = CLASS_NAMES[np.argmax(preds)]
    confidence = round(float(np.max(preds)), 3)
    return f"{label} ({confidence})"

# ✅ Step 5: Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Radiograph"),
    outputs=gr.Text(label="Prediction"),
    title="🦴 OsteoLogic Fracture Detector",
    description="Upload a radiograph to detect fractures using safetensors-powered MobileNetV2."
).launch()