import sys
import os
import cv2
import torch
import numpy as np
import time

sys.path.insert(0, 'yolov7')

from picamera2 import Picamera2
from models.experimental import attempt_load
from utils.general       import non_max_suppression, scale_coords
from utils.torch_utils   import select_device
from utils.datasets      import letterbox

# ─── CONFIGURATION ───────────────────────────────────────────────

WEIGHTS    = "best.pt"
IMG_SIZE   = 640
CONF_THRES = 0.40
IOU_THRES  = 0.45

CLASSES = [
    "MMs_peanut", "MMs_regular", "airheads",
    "gummy_worms", "milky_way", "nerds",
    "skittles", "snickers", "starbust",
    "three_musketeers", "twizzlers"
]

COLORS = [
    (255,  60,  60), (255, 160,  30), ( 50, 205,  50),
    (255,  20, 147), (138,  43, 226), ( 30, 144, 255),
    (255, 215,   0), (220,  20,  60), (255, 105, 180),
    ( 64, 224, 208), (255, 140,   0),
]

# ─── LOAD MODEL ──────────────────────────────────────────────────

print("Loading YOLOv7 model...")
device = select_device('')   # CPU on Pi
model  = attempt_load(WEIGHTS, map_location=device)
model.eval()
stride = int(model.stride.max())
print("Model loaded!")

# ─── START PI CAMERA ─────────────────────────────────────────────

print("Starting Pi Camera...")
picam = Picamera2()
config = picam.create_preview_configuration(
    main={"size": (640, 480), "format": "RGB888"}
)
picam.configure(config)
picam.start()
time.sleep(2)   # camera warm-up
print("Camera started!")

print("\nRunning detection... Press Q to quit")
print("-" * 60)

frame_count = 0
fps_start    = time.time()

while True:
    # ── Capture frame from Pi Camera ─────────────────────────────
    frame = picam.capture_array()


    orig_frame = frame.copy()
    h0, w0     = frame.shape[:2]

    # ── Preprocess for YOLOv7 ────────────────────────────────────
    img = letterbox(frame, IMG_SIZE, stride=stride)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device).float() / 255.0
    img = img.unsqueeze(0)

    # ── Inference ─────────────────────────────────────────────────
    with torch.no_grad():
        pred = model(img)[0]
        pred = non_max_suppression(pred, CONF_THRES, IOU_THRES)

    # ── Draw detections ───────────────────────────────────────────
    detections = []
    if pred[0] is not None and len(pred[0]):
        det = pred[0].clone()
        scale_coords(img.shape[2:], det[:, :4], (h0, w0)).round()

        for *xyxy, conf, cls in det:
            x1, y1, x2, y2 = map(int, xyxy)
            cls_id = int(cls)
            label  = CLASSES[cls_id] if cls_id < len(CLASSES) else "unknown"
            color  = COLORS[cls_id % len(COLORS)]

            cx = (x1 + x2) // 2
            cy = (y1 + y2) // 2

            detections.append({
                "class": label, "conf": float(conf),
                "pixel": (cx, cy), "bbox": (x1, y1, x2, y2)
            })

            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            text = f"{label} {conf:.2f}"
            cv2.putText(frame, text, (x1, y1 - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            cv2.circle(frame, (cx, cy), 4, color, -1)

            print(f"  {label:20s} pixel=({cx},{cy}) conf={conf:.2f}")

    # ── FPS counter ────────────────────────────────────────────────
    frame_count += 1
    if frame_count % 30 == 0:
        fps = 30 / (time.time() - fps_start)
        fps_start = time.time()
        print(f"  FPS: {fps:.1f}")

    cv2.putText(frame, f"Detections: {len(detections)}",
                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)

    # ── Show frame ──────────────────────────────────────────────────
    cv2.imshow("Pi Camera - Candy Detection", frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

picam.stop()
cv2.destroyAllWindows()
print("Stopped.")