import cv2
import numpy as np
import os
import time

# تنظیمات اصلی
USE_PHONE_CAMERA = False  # True اگر از گوشی به عنوان وب‌کم استفاده می‌کنید
CONFIDENCE_THRESHOLD = 0.4  # تنظیم برای تشخیص سریع‌تر
NMS_THRESHOLD = 0.2         # تنظیم برای کاهش تداخل جعبه‌ها
HISTORY_DURATION = 0.1        # مدت زمان نگهداری تشخیص‌ها (ثانیه)
MOTION_THRESHOLD = 30       # آستانه تشخیص حرکت
INPUT_SIZE = 224            # سایز ورودی مدل برای افزایش FPS

# مدل و وزن‌های از پیش آموزش دیده
MODEL_WEIGHTS = "yolov4-tiny.weights"
MODEL_CONFIG = "yolov4-tiny.cfg"
CLASS_NAMES = "coco.names"

# کلاس‌های مورد نظر: پرنده، ماشین، انسان
TARGET_CLASSES = ["bird", "car", "person"]

# بررسی وجود فایل‌های ضروری
for file_path in [MODEL_WEIGHTS, MODEL_CONFIG, CLASS_NAMES]:
    if not os.path.exists(file_path):
        print(f"خطا: فایل {file_path} یافت نشد!")
        exit()

# بارگیری مدل
print("بارگیری مدل...")
net = cv2.dnn.readNet(MODEL_WEIGHTS, MODEL_CONFIG)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)  # استفاده از CPU

# بارگیری کلاس‌ها
with open(CLASS_NAMES, "r") as f:
    classes = [line.strip() for line in f.readlines()]
    
# تعیین کلاس‌های مورد نظر
class_ids = [i for i, name in enumerate(classes) if name in TARGET_CLASSES]

# تعیین لایه‌های خروجی
try:
    output_layers = net.getUnconnectedOutLayersNames()
except:
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# راه‌اندازی دوربین
print("راه‌اندازی دوربین...")
if USE_PHONE_CAMERA:
    PHONE_IP = "192.168.1.100"  # آدرس IP گوشی خود را وارد کنید
    PHONE_PORT = "4747"         # پورت DroidCam
    camera_url = f"http://{PHONE_IP}:{PHONE_PORT}/video"
    cap = cv2.VideoCapture(camera_url)
else:
    cap = cv2.VideoCapture(0)  # استفاده از وب‌کم لپ‌تاپ

if not cap.isOpened():
    print("خطا: دوربین باز نشد!")
    exit()

# تنظیم کیفیت دوربین برای افزایش FPS
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_FPS, 60)

print("سیستم تشخیص اشیا فعال شد!")
print("تشخیص: پرنده، ماشین، انسان")
print("برای خروج کلید 'Q' را فشار دهید")

# رنگ‌ها برای کلاس‌های مختلف
colors = {
    "bird": (0, 255, 255),    # زرد
    "car": (255, 0, 0),       # آبی
    "person": (0, 0, 255)     # قرمز
}

# لیست تشخیص‌های اخیر برای نگهداری 5 ثانیه‌ای
detection_history = []

# متغیرهای تشخیص حرکت
prev_frame = None
motion_mask = None
prev_time = time.time()

# متغیرهای FPS
fps = 0
frame_count = 0
start_time = time.time()

# ایجاد پنجره‌ها
cv2.namedWindow('Object Detection', cv2.WINDOW_NORMAL)
cv2.namedWindow('Motion Detection', cv2.WINDOW_NORMAL)

try:
    while True:
        # دریافت فریم
        ret, frame = cap.read()
        if not ret:
            print("خطا در دریافت فریم!")
            time.sleep(0.1)
            continue
        
        # افزایش شمارنده فریم‌ها
        frame_count += 1
        
        # محاسبه FPS هر ثانیه
        current_time = time.time()
        elapsed = current_time - start_time
        if elapsed >= 1.0:
            fps = frame_count / elapsed
            frame_count = 0
            start_time = current_time
            print(f"FPS: {fps:.1f}")
        
        # کاهش سایز فریم برای پردازش سریع‌تر
        small_frame = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
        height, width = frame.shape[:2]
        
        # --- تشخیص حرکت ---
        gray = cv2.cvtColor(small_frame, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (21, 21), 0)
        
        if prev_frame is None:
            prev_frame = gray
            continue
        
        # محاسبه تفاضل فریم‌ها
        frame_delta = cv2.absdiff(prev_frame, gray)
        thresh = cv2.threshold(frame_delta, MOTION_THRESHOLD, 255, cv2.THRESH_BINARY)[1]
        thresh = cv2.dilate(thresh, None, iterations=2)
        
        # ذخیره فریم فعلی برای فریم بعدی
        prev_frame = gray
        
        # ایجاد ماسک حرکت برای نمایش
        motion_display = cv2.cvtColor(thresh, cv2.COLOR_GRAY2BGR)
        motion_display = cv2.resize(motion_display, (width, height))
        
        # --- تشخیص اشیاء با YOLO ---
        blob = cv2.dnn.blobFromImage(small_frame, 0.00392, (INPUT_SIZE, INPUT_SIZE), swapRB=True, crop=False)
        net.setInput(blob)
        outs = net.forward(output_layers)
        
        # پردازش نتایج
        new_detections = []
        
        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                
                # فیلتر کلاس‌های مورد نظر و اطمینان
                if confidence > CONFIDENCE_THRESHOLD and class_id in class_ids:
                    # محاسبه مختصات جعبه در فریم اصلی
                    center_x = int(detection[0] * width)
                    center_y = int(detection[1] * height)
                    w = int(detection[2] * width)
                    h = int(detection[3] * height)
                    
                    # مختصات گوشه بالا-چپ
                    x = max(0, int(center_x - w / 2))
                    y = max(0, int(center_y - h / 2))
                    
                    # فیلتر اشیای خیلی کوچک
                    if w > 10 and h > 10:
                        new_detections.append({
                            'class_id': class_id,
                            'confidence': float(confidence),
                            'box': [x, y, w, h],
                            'timestamp': current_time
                        })
        
        # حذف جعبه‌های تکراری با NMS
        boxes = [det['box'] for det in new_detections]
        confidences = [det['confidence'] for det in new_detections]
        
        if len(boxes) > 0:
            indices = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)
            if len(indices) > 0:
                # تبدیل اندیس‌ها به لیست تخت
                if hasattr(indices, 'shape'):
                    indices = indices.flatten()
                else:
                    indices = [i[0] for i in indices]
                
                new_detections = [new_detections[i] for i in indices]
        
        # اضافه کردن تشخیص‌های جدید به تاریخچه
        detection_history.extend(new_detections)
        
        # حذف تشخیص‌های قدیمی (بیشتر از 5 ثانیه)
        detection_history = [det for det in detection_history 
                            if current_time - det['timestamp'] <= HISTORY_DURATION]
        
        # --- رسم نتایج روی فریم اصلی ---
        # رسم تمام تشخیص‌های فعال در تاریخچه
        for det in detection_history:
            x, y, w, h = det['box']
            class_id = det['class_id']
            class_name = classes[class_id]
            confidence = det['confidence']
            color = colors.get(class_name, (0, 255, 0))
            
            # محاسبه شفافیت بر اساس زمان باقیمانده
            time_left = HISTORY_DURATION - (current_time - det['timestamp'])
            alpha = max(0.3, time_left / HISTORY_DURATION)  # شفافیت از 0.3 تا 1.0
            
            # رسم جعبه و برچسب
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            
            label = f"{class_name} {confidence:.2f}"
            cv2.putText(frame, label, (x, y - 10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
            
            # افزودن نقطه مرکزی برای ردیابی بهتر
            center = (x + w//2, y + h//2)
            cv2.circle(frame, center, 5, color, -1)
        
        # --- بهبود تشخیص با استفاده از اطلاعات حرکت ---
        # ترکیب اطلاعات حرکت با تشخیص‌های YOLO
        for det in detection_history:
            x, y, w, h = det['box']
            
            # استخراج ناحیه مربوطه در ماسک حرکت
            motion_roi = motion_display[y:y+h, x:x+w]
            
            # محاسبه درصد حرکت در ناحیه
            if motion_roi.size > 0:
                motion_percentage = np.mean(motion_roi) / 255
                
                # اگر حرکت قابل توجهی وجود دارد، جعبه را برجسته‌تر نمایش دهیم
                if motion_percentage > 0.2:
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 4)
                    cv2.putText(frame, "MOVING!", (x, y - 30), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        # --- نمایش اطلاعات ---
        cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(frame, "Detection: Bird, Car, Person", (10, 60), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(frame, f"Tracking: {len(detection_history)} objects", (10, 90), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        # --- نمایش خروجی‌ها ---
        # نمایش فریم اصلی با تشخیص‌ها
        cv2.imshow('Object Detection', frame)
        
        # نمایش ماسک حرکت با اطلاعات اضافه
        cv2.putText(motion_display, "MOTION DETECTION", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(motion_display, "White: Fast Motion", (10, 60), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
        cv2.putText(motion_display, "Gray: Slow Motion", (10, 90), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (150, 150, 150), 1)
        cv2.putText(motion_display, "Black: No Motion", (10, 120), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
        cv2.imshow('Motion Detection', motion_display)
        
        # خروج با فشار Q
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

except Exception as e:
    print(f"خطای غیرمنتظره: {e}")

finally:
    cap.release()
    cv2.destroyAllWindows()
    print("منابع سیستم آزاد شدند")
