import os, json
import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import preprocess_input

# -----------------------------
# Environment
# -----------------------------
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

# -----------------------------
# Load model
# -----------------------------
MODEL_PATH = "solarguard_resnet_final.keras"
CLASS_JSON = "class_indices.json"

model = tf.keras.models.load_model(MODEL_PATH, compile=False)

# -----------------------------
# Load class mapping (CRITICAL)
# -----------------------------
with open(CLASS_JSON, "r") as f:
    class_indices = json.load(f)

# reverse index -> class
idx_to_class = {v: k for k, v in class_indices.items()}

IMG_SIZE = (224, 224)

# -----------------------------
# Preprocess (same as training)
# -----------------------------
def preprocess_image(img):
    img = img.resize(IMG_SIZE)
    arr = np.array(img).astype("float32")
    arr = np.expand_dims(arr, axis=0)
    return preprocess_input(arr)

# -----------------------------
# Predict
# -----------------------------
def predict(img: Image.Image):
    img = img.convert("RGB")
    x = preprocess_image(img)

    preds = model(x, training=False).numpy()[0]
    idx = int(np.argmax(preds))

    return {
        "Prediction": idx_to_class[idx],
        "Confidence": f"{preds[idx] * 100:.2f}%",
        "Probabilities": {
            idx_to_class[i]: float(preds[i])
            for i in range(len(preds))
        }
    }

# -----------------------------
# UI
# -----------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Solar Panel Image"),
    outputs=gr.JSON(label="Result"),
    title="⚡ SolarGuard – Solar Panel Defect Detection",
    description="Accurate ResNet50-based solar panel defect classification"
)

if __name__ == "__main__":
    demo.launch()