import os
import time
import numpy as np
from PIL import Image
import warnings
warnings.filterwarnings("ignore")

import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.vgg16        import preprocess_input as vgg_preprocess
from tensorflow.keras.applications.resnet50     import preprocess_input as resnet_preprocess
from tensorflow.keras.applications.mobilenet_v2  import preprocess_input as mobilenet_preprocess
from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

IMG_SIZE   = 224
MODEL_REPO = "GeetamSharma/smartvision-models"

IDX_TO_CLASS = {
    0:"airplane",    1:"bed",          2:"bench",
    3:"bicycle",     4:"bird",         5:"bottle",
    6:"bowl",        7:"bus",          8:"cake",
    9:"car",        10:"cat",         11:"chair",
   12:"couch",      13:"cow",         14:"cup",
   15:"dog",        16:"elephant",    17:"horse",
   18:"motorcycle", 19:"person",      20:"pizza",
   21:"potted_plant",22:"stop_sign",  23:"traffic_light",
   24:"truck"
}

def download_model(filename):
    path = hf_hub_download(
        repo_id   = MODEL_REPO,
        filename  = filename,
        repo_type = "model",
    )
    return path

def load_all_models(models_dir=None):
    models = {}
    print("Loading VGG16...")
    models["vgg16"]        = load_model(
        download_model("vgg16_best.keras"), compile=False)
    print("Loading ResNet50...")
    models["resnet50"]     = load_model(
        download_model("resnet50_best.keras"), compile=False)
    print("Loading MobileNetV2...")
    models["mobilenet"]    = load_model(
        download_model("mobilenet_best.keras"), compile=False)
    print("Loading EfficientNetB0...")
    models["efficientnet"] = load_model(
        download_model("efficientnet_best.keras"), compile=False)
    print("Loading YOLOv8...")
    models["yolo"]         = YOLO(download_model("yolov8_best.pt"))
    print("All models loaded!")
    return models

def preprocess_image(image, preprocess_fn):
    if image.mode != "RGB":
        image = image.convert("RGB")
    img = image.resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS)
    arr = np.array(img, dtype=np.float32)
    arr = np.expand_dims(arr, axis=0)
    arr = preprocess_fn(arr)
    return arr

def classify_image(image, models, top_k=5):
    models_config = [
        ("VGG16",         models["vgg16"],        vgg_preprocess),
        ("ResNet50",       models["resnet50"],      resnet_preprocess),
        ("MobileNetV2",    models["mobilenet"],     mobilenet_preprocess),
        ("EfficientNetB0", models["efficientnet"],  efficientnet_preprocess),
    ]
    all_predictions = {}
    for model_name, model, preprocess_fn in models_config:
        start    = time.time()
        arr      = preprocess_image(image, preprocess_fn)
        preds    = model.predict(arr, verbose=0)[0]
        top_idx  = np.argsort(preds)[::-1][:top_k]
        top_data = []
        for idx in top_idx:
            top_data.append({
                "class"     : IDX_TO_CLASS[int(idx)],
                "confidence": float(preds[idx]),
                "percentage": float(preds[idx] * 100),
            })
        elapsed = (time.time() - start) * 1000
        all_predictions[model_name] = {
            "top_prediction": top_data[0]["class"],
            "top_confidence": top_data[0]["confidence"],
            "top_k"         : top_data,
            "inference_ms"  : round(elapsed, 2),
        }
    return all_predictions

def detect_objects(image, models,
                   conf_threshold=0.25, iou_threshold=0.45):
    if image.mode != "RGB":
        image = image.convert("RGB")
    start   = time.time()
    results = models["yolo"].predict(
        source=image, conf=conf_threshold,
        iou=iou_threshold, imgsz=640, verbose=False,
    )
    elapsed    = (time.time() - start) * 1000
    result     = results[0]
    detections = []
    if result.boxes is not None and len(result.boxes) > 0:
        for box in result.boxes:
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
            confidence      = float(box.conf[0].cpu().numpy())
            class_idx       = int(box.cls[0].cpu().numpy())
            detections.append({
                "class_name": IDX_TO_CLASS[class_idx],
                "confidence": confidence,
                "percentage": round(confidence * 100, 2),
                "bbox"      : [float(x1), float(y1), float(x2), float(y2)],
                "class_idx" : class_idx,
            })
    detections    = sorted(detections, key=lambda x: x["confidence"], reverse=True)
    annotated_pil = Image.fromarray(result.plot()[:, :, ::-1])
    return {
        "detections"     : detections,
        "num_detections" : len(detections),
        "annotated_image": annotated_pil,
        "inference_ms"   : round(elapsed, 2),
    }