import spaces
import gradio as gr
import torch
import os
import sys
import glob
import cv2
import numpy as np

# Skip torch hub trust check for HF Spaces deployment
os.environ['TORCH_HOME'] = '/tmp/torch_home'

# --- Detection model (finds objects + bounding boxes in the scene) ---
# 'custom' + path= tells torch.hub to load your own YOLOv5-trained weights
# via the YOLOv5 repo code (this is what handles the old-style checkpoint format).
model = torch.hub.load("ultralytics/yolov5", "custom", path="best.pt", trust_repo=True)

# Loading the detection model above causes torch.hub to clone/cache the yolov5
# repo and add it to sys.path so its 'models' package is importable. We rely on
# that same sys.path entry below to unpickle the classifier checkpoint, but add
# a defensive fallback in case caching behavior ever changes.
_hub_repo_dir = os.path.join(torch.hub.get_dir(), "ultralytics_yolov5_master")
if os.path.isdir(_hub_repo_dir) and _hub_repo_dir not in sys.path:
    sys.path.insert(0, _hub_repo_dir)

# --- Classification model (identifies the specific airplane type from a crop) ---
CLASSIFIER_WEIGHTS = "best_aeroplane.pt"
CLS_IMGSZ = 224
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

classifier_model = None
classifier_names = None

def load_classifier():
    """Loads the native YOLOv5 classification checkpoint used to identify airplane sub-types."""
    global classifier_model, classifier_names
    if not os.path.isfile(CLASSIFIER_WEIGHTS):
        print(f"WARNING: classifier weights '{CLASSIFIER_WEIGHTS}' not found; "
              f"airplane detections will fall back to the generic 'airplane' label.")
        return
    # weights_only=False: required on PyTorch >=2.6 since YOLOv5 checkpoints pickle
    # full model objects (models.yolo.ClassificationModel). Safe for a checkpoint you trained.
    ckpt = torch.load(CLASSIFIER_WEIGHTS, map_location="cpu", weights_only=False)
    classifier_model = (ckpt.get("ema") or ckpt["model"]).float().eval()
    classifier_names = classifier_model.names


load_classifier()


def classify_crop(crop_rgb):
    """Runs the airplane sub-type classifier on a cropped RGB image region.

    Returns (label, confidence), or (None, 0.0) if the classifier isn't loaded
    or the crop is empty.
    """
    if classifier_model is None or crop_rgb is None or crop_rgb.size == 0:
        return None, 0.0

    resized = cv2.resize(crop_rgb, (CLS_IMGSZ, CLS_IMGSZ), interpolation=cv2.INTER_LINEAR)
    img = resized.astype(np.float32) / 255.0
    img = (img - IMAGENET_MEAN) / IMAGENET_STD
    tensor = torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0)

    with torch.no_grad():
        logits = classifier_model(tensor)
        probs = torch.nn.functional.softmax(logits, dim=1)
        conf, idx = probs.max(dim=1)

    label = classifier_names[int(idx.item())]
    return label, float(conf.item())


font_size = 1.0
font_thickness = 1
box_width = 1

SAMPLE_IMAGES_DIR = "Images"
SAMPLE_IMAGE_EXTS = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")


def get_sample_images():
    """Returns a sorted list of image file paths found in the Images folder."""
    files = []
    for ext in SAMPLE_IMAGE_EXTS:
        files.extend(glob.glob(os.path.join(SAMPLE_IMAGES_DIR, ext)))
        files.extend(glob.glob(os.path.join(SAMPLE_IMAGES_DIR, ext.upper())))
    return sorted(set(files))


def show_sample_gallery():
    images = get_sample_images()
    if not images:
        gr.Warning(f"No images found in the '{SAMPLE_IMAGES_DIR}' folder.")
        return gr.update(visible=False, value=[], columns=8, height=250,
                          object_fit="cover", buttons=["fullscreen"]), []
    return gr.update(visible=True, value=images, columns=8, height=250,
                      object_fit="cover", buttons=["fullscreen"]), images


def select_sample_image(images_list, evt: gr.SelectData):
    selected = None
    if images_list and 0 <= evt.index < len(images_list):
        selected = images_list[evt.index]
    return selected, gr.update(visible=False, columns=8, height=250,
                                object_fit="cover", buttons=["fullscreen"])
                                

def annotate_with_custom_font(image, results):
    global font_size
    global font_thickness
    global box_width

    """Annotate image with custom font size, thickness, and bounding box width.
    Airplane detections are re-labeled using the airplane sub-type classifier."""
    # Get image dimensions
    if isinstance(image, np.ndarray):
        h, w = image.shape[:2]
        img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    else:
        img_bgr = image

    # Draw bounding boxes and text with custom settings
    for det in results.xyxy[0]:
        x1, y1, x2, y2, conf, cls = det.cpu().numpy()
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

        class_name = model.names[int(cls)]
        label = f"{class_name} {conf:.2f}"

        # If the detector found an airplane, crop it out and hand it to the
        # classifier to get the specific airplane type instead of the generic label.
        if class_name.lower() in ("airplane", "aeroplane"):
            x1c, y1c = max(0, x1), max(0, y1)
            x2c, y2c = max(x1c, x2), max(y1c, y2)
            crop_rgb = image[y1c:y2c, x1c:x2c]
            cls_label, cls_conf = classify_crop(crop_rgb)
            if cls_label is not None:
                label = f"{cls_label} {cls_conf:.2f}"

        # Draw bounding box with custom width
        cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (0, 255, 0), box_width)

        # Draw text with custom font settings
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = font_size
        text_thickness = font_thickness

        text_size = cv2.getTextSize(label, font, font_scale, text_thickness)[0]
        text_x = x1
        text_y = max(y1 - 10, text_size[1] + 5)

        cv2.putText(img_bgr, label, (text_x, text_y), font,
                    font_scale, (0, 255, 0), text_thickness)

    # Convert back to RGB
    annotated_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    return annotated_rgb


@spaces.GPU
def process_image(image):
    if image is None:
        return None
    model.to("cpu")
    results = model(image)
    annotated = annotate_with_custom_font(image, results)
    return annotated


def load_sample_video():
    return "video.mp4"


@spaces.GPU
def process_video(video_path):
    if video_path is None:
        return None

    model.to("cpu")

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 24
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Collect frames first
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)
        annotated_frame = annotate_with_custom_font(frame, results)
        frames.append(annotated_frame)

    cap.release()

    # Try imageio first (better codec support), fallback to cv2
    output_path = "output.mp4"
    try:
        import imageio
        writer = imageio.get_writer(output_path, fps=fps, codec='libx264')
        for frame in frames:
            # Convert RGB to BGR for imageio (imageio expects BGR like OpenCV)
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            writer.append_data(frame_bgr)
        writer.close()
    except (ImportError, Exception) as e:
        print(f"Imageio failed: {e}, falling back to OpenCV")
        # Fallback to OpenCV with MPEG-4 codec (better compatibility)
        fourcc = cv2.VideoWriter_fourcc(*"MPEG")
        writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        for frame in frames:
            # frame is RGB from render() - convert to BGR for OpenCV
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            writer.write(frame_bgr)
        writer.release()

    return output_path


with gr.Blocks() as demo:
    gr.Markdown("# YOLO Object Detection")
    gr.Markdown("Upload an image or a video, then click **Process** to run YOLO detection.")

    with gr.Tabs():
        with gr.Tab("Image"):
            sample_gallery = gr.Gallery(
                label="Choose a sample image",
                visible=False,
                preview=False,
                allow_preview=False,
                columns=8,
                height=250,
                object_fit="cover",
                buttons=["fullscreen"],
            )
            sample_images_state = gr.State([])

            with gr.Row():
                load_sample_btn = gr.Button("Load Sample Image", variant="secondary")
            load_sample_btn.click(
                fn=show_sample_gallery,
                outputs=[sample_gallery, sample_images_state],
            )
            with gr.Column(scale=1):
                image_input = gr.Image(type="numpy", label="Input Image")
            sample_gallery.select(
                fn=select_sample_image,
                inputs=[sample_images_state],
                outputs=[image_input, sample_gallery],
            )
            with gr.Row():
                image_button = gr.Button("Process", variant="primary")
            with gr.Column(scale=1):
                image_output = gr.Image(type="numpy", label="Detections")

            image_button.click(fn=process_image, inputs=[image_input], outputs=image_output)

        with gr.Tab("Video"):
            with gr.Row():
                load_sample_video_btn = gr.Button("Load Sample Video", variant="secondary")
            with gr.Column(scale=1):
                video_input = gr.Video(label="Input Video")
            with gr.Row():
                video_button = gr.Button("Process", variant="primary")
            with gr.Column(scale=1):
                video_output = gr.Video(label="Detections")
            load_sample_video_btn.click(fn=load_sample_video, outputs=video_input)
            video_button.click(fn=process_video, inputs=[video_input], outputs=video_output)

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