import torch
from transformers import DetrFeatureExtractor, DetrForObjectDetection
from transformers import ViTImageProcessor, ViTModel
from PIL import Image, ImageDraw, ImageFont
import requests
import numpy as np
from sklearn.ensemble import IsolationForest
import gradio as gr

# 1. Load pretrained models
detr_feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
detr_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')

vit_processor = ViTImageProcessor.from_pretrained('facebook/dino-vitb16')
vit_model = ViTModel.from_pretrained('facebook/dino-vitb16')

# COCO classes for DETR object detection
COCO_CLASSES = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
    'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]

# 2. Function: Object detection with DETR
def detect_objects(image, confidence_threshold=0.7):
    inputs = detr_feature_extractor(images=image, return_tensors="pt")
    outputs = detr_model(**inputs)

    logits = outputs.logits
    bboxes = outputs.pred_boxes

    probas = logits.softmax(-1)[0, :, :-1]  # ignore last class (no object)
    keep = probas.max(-1).values > confidence_threshold
    scores, labels = probas.max(-1)
    boxes = bboxes[0, keep]

    detected_objects = []
    for score, label_idx, box in zip(scores[keep], labels[keep], boxes):
        box = box.detach().cpu().numpy()
        # Convert box from [center_x, center_y, width, height] normalized format to pixel format
        w, h = image.size
        cx, cy, bw, bh = box
        x_min = int((cx - bw / 2) * w)
        y_min = int((cy - bh / 2) * h)
        x_max = int((cx + bw / 2) * w)
        y_max = int((cy + bh / 2) * h)

        detected_objects.append({
            'label': COCO_CLASSES[label_idx],
            'score': float(score),
            'box': [x_min, y_min, x_max, y_max]
        })

    return detected_objects

# 3. Function: Scene Recognition with ViT embeddings (dummy classifier)
def get_scene_embedding(image):
    inputs = vit_processor(images=image, return_tensors="pt")
    outputs = vit_model(**inputs)
    cls_embedding = outputs.last_hidden_state[:, 0, :].detach().cpu().numpy()  # shape (1, hidden_dim)
    return cls_embedding

# For demonstration: A simple heuristic scene classifier with hardcoded labels
def classify_scene(embedding):
    # In practice, train a scene classifier on embeddings. Here dummy logic:
    avg_value = np.mean(embedding)
    if avg_value > 0:
        return "Indoor Scene"
    else:
        return "Outdoor Scene"

# 4. Simple anomaly detection combining object detection confidence + scene embedding
# For demo, anomaly if no object detected OR average detection score below threshold OR unusual scene embedding

def is_anomalous(detected_objects, scene_embedding):
    if len(detected_objects) == 0:
        return True, "No objects detected - possible anomaly"

    avg_score = np.mean([obj['score'] for obj in detected_objects])
    if avg_score < 0.75:
        return True, "Low object detection confidence - possible anomaly"

    # Simple heuristic on embedding mean (you can replace with IsolationForest or more)
    if np.mean(scene_embedding) < -0.1:
        return True, "Unusual scene context detected"

    return False, "Normal"

# 5. Draw bounding boxes and labels on image
def draw_detections(image, detected_objects, anomaly_flag, anomaly_reason, scene_label):
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()

    for obj in detected_objects:
        box = obj['box']
        label = obj['label']
        score = obj['score']
        draw.rectangle(box, outline="red", width=3)
        draw.text((box[0], box[1] - 10), f"{label}: {score:.2f}", fill="red", font=font)

    # Draw anomaly status and scene label
    draw.text((10, 10), f"Scene: {scene_label}", fill="yellow", font=font)
    anomaly_text = f"Anomaly Detected: {anomaly_flag} - {anomaly_reason}"
    draw.text((10, 30), anomaly_text, fill="yellow", font=font)

    return image

# 6. Main processing function for Gradio app
def process_image(img):
    detected_objects = detect_objects(img)
    scene_emb = get_scene_embedding(img)
    scene_label = classify_scene(scene_emb)
    anomaly_flag, anomaly_reason = is_anomalous(detected_objects, scene_emb)

    img_with_boxes = img.copy()
    img_with_boxes = draw_detections(img_with_boxes, detected_objects, anomaly_flag, anomaly_reason, scene_label)
    return img_with_boxes, anomaly_flag, anomaly_reason, scene_label

# 7. Build Gradio Interface for Hugging Face Spaces
title = "Advanced Vision Pipeline: Object Detection + Scene Recognition + Anomaly Detection"

iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Output Image with Detections"),
        gr.Textbox(label="Anomaly Detected?"),
        gr.Textbox(label="Anomaly Reason"),
        gr.Textbox(label="Scene Classification"),
    ],
    title=title,
    description="Upload an image to detect objects, recognize scene context, and detect unusual anomalies."
)

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