from __future__ import annotations

import io
import os
from pathlib import Path
from typing import Dict, List

import gradio as gr
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow.keras.applications import EfficientNetB2


# Label wajah sesuai backend Flask
CLASS_NAMES: List[str] = ["Heart", "Oblong", "Oval", "Round", "Square"]
IMG_SIZE = 244
NUM_CLASSES = len(CLASS_NAMES)


def _load_image_to_rgb(image: Image.Image) -> np.ndarray:
    if image.mode != "RGB":
        image = image.convert("RGB")
    return np.asarray(image)


def _resize_to_model(img_rgb: np.ndarray) -> np.ndarray:
    im = Image.fromarray(img_rgb)
    im = im.resize((IMG_SIZE, IMG_SIZE), Image.NEAREST)
    return np.asarray(im)


def _preprocess(image: Image.Image) -> np.ndarray:
    rgb = _load_image_to_rgb(image)
    resized = _resize_to_model(rgb)
    arr = resized.astype("float32")
    return np.expand_dims(arr, axis=0)  # [1, IMG_SIZE, IMG_SIZE, 3]


def _create_model(num_classes: int):
    base_model = EfficientNetB2(
        weights=None,
        include_top=False,
        pooling="avg",
        input_shape=(IMG_SIZE, IMG_SIZE, 3),
    )
    model = tf.keras.Sequential(
        [
            base_model,
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(num_classes, activation="softmax"),
        ]
    )
    return model


def _resolve_model_path() -> Path:
    base_dir = Path(__file__).resolve().parent
    candidates = [
        base_dir / "best_model.keras",
    ]
    for p in candidates:
        if p.exists():
            return p
    # default fallback
    return candidates[0]


class PreTrainedModel:
    def __init__(self) -> None:
        self.model_path = _resolve_model_path()
        self.model = _create_model(NUM_CLASSES)
        if self.model_path.exists():
            # Memuat bobot yang sama seperti backend Flask
            self.model.load_weights(self.model_path)
            print(f"Loaded weights from: {self.model_path}")
        else:
            print(f"Warning: Model file not found at: {self.model_path}")

    def predict_image(self, image: Image.Image) -> Dict[str, float]:
        x = _preprocess(image)
        preds = self.model.predict(x)
        if isinstance(preds, (list, tuple)):
            preds = preds[0]
        probs = np.asarray(preds).squeeze().tolist()
        return {label: float(score) for label, score in zip(CLASS_NAMES, probs)}


model = PreTrainedModel()


def predict(image: Image.Image):
    predictions = model.predict_image(image)
    probs_percent = {label: round(p * 100, 2)
                     for label, p in predictions.items()}
    max_label = max(probs_percent, key=probs_percent.get)
    return {
        "label": max_label,
        "percentage": probs_percent[max_label],
        "probabilities": probs_percent,
    }


iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="Face Shape Classification",
    description="Upload an image to classify face shape (Heart, Oblong, Oval, Round, Square).",
)


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