import os
import zipfile
import tensorflow as tf
import numpy as np
import gradio as gr
import gdown
from tensorflow.keras.preprocessing import image as keras_image
from PIL import Image

# === Google Drive ZIP file ID ===
ZIP_FILE_ID = "1KYs5K2cIKp6C5VlIkATjoAJnUyDUhaFx"  # from your shared link
ZIP_PATH = "best_model (3).keras.zip"
MODEL_PATH = "model.weights.h5"  # expected inside ZIP after extraction

def download_and_extract():
    if not os.path.exists(MODEL_PATH):
        print("Downloading model ZIP from Google Drive...")
        url = f"https://drive.google.com/uc?id={ZIP_FILE_ID}"
        gdown.download(url, ZIP_PATH, quiet=False)
        print("Extracting ZIP...")
        with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
            zip_ref.extractall()
        print("Extraction complete.")

# === Download and extract the model on app startup ===
download_and_extract()

# === Load the model ===
model = tf.keras.models.load_model(MODEL_PATH)

# === Class names — adjust based on your model's labels ===
class_names = ['Dry', 'Normal', 'Oily', 'Acne', 'Blackheads', 'Dark Spots', 'Wrinkles', 'Skin Redness', 'Pores', 'Eye Bags']

def preprocess_image(img):
    img = img.convert("RGB")
    img = img.resize((224, 224))  # adjust size if your model expects 225×225
    arr = keras_image.img_to_array(img) / 255.0
    return np.expand_dims(arr, axis=0)

def predict(img):
    inp = preprocess_image(img)
    preds = model.predict(inp, verbose=0)[0]
    top_idx = np.argmax(preds)
    top_pred = f"{class_names[top_idx]} ({preds[top_idx]*100:.2f}%)"
    all_probs = "\n".join(f"{class_names[i]}: {v*100:.2f}%" for i, v in enumerate(preds))
    return top_pred, all_probs

# === Gradio Interface ===
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Textbox(label="Top Prediction"), gr.Textbox(label="All Class Probabilities")],
    title="Skin Type & Condition Predictor"
)

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