import cv2
import numpy as np
import pymysql
from datetime import datetime

# Database connection
conn = pymysql.connect(
    host='localhost',
    user='root',
    password='',
    database='db_traffic2'
)
cursor = conn.cursor()

# RTSP stream or test video
cap = cv2.VideoCapture('pagi.mp4')  # Replace with RTSP URL when needed

# Get original width and height
ret, frame = cap.read()
if not ret:
    print("Failed to read from video source.")
    exit()

height2, width2, channels = frame.shape

# Background subtractor
fgbg = cv2.createBackgroundSubtractorMOG2()

# Vehicle tracking and speed
vehicle_id_counter = 0
tracker_dict = {}
speed_line_y1 = 245
speed_line_y2 = 355
meters_travel_pixel = 0.2

last_vehicle_ids_in_zone = set()

while True:
    ret, frame = cap.read()
    if not ret:
        break

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    fgmask = fgbg.apply(gray)

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    closing = cv2.morphologyEx(fgmask, cv2.MORPH_CLOSE, kernel)
    opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)
    dilation = cv2.dilate(opening, kernel)
    _, bins = cv2.threshold(dilation, 220, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(bins, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
    hull = [cv2.convexHull(c) for c in contours]
    cv2.drawContours(frame, hull, -1, (0, 255, 0), 2)

    min_area = 400
    max_area = 40000

    vehicle_ids_in_zone = set()

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if min_area < area < max_area:
            M = cv2.moments(cnt)
            if M['m00'] == 0:
                continue
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])

            if speed_line_y1 <= cy <= speed_line_y2:
                vehicle_ids_in_zone.add(matched_id)

            x, y, w, h = cv2.boundingRect(cnt)
            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
            cv2.drawMarker(frame, (cx, cy), (0, 0, 255), cv2.MARKER_CROSS, 10, 1)

            matched_id = None
            for vid, info in tracker_dict.items():
                old_cx, old_cy = info['pos']
                if abs(cx - old_cx) < 30 and abs(cy - old_cy) < 30:
                    matched_id = vid
                    break

            if matched_id is None:
                matched_id = vehicle_id_counter
                vehicle_id_counter += 1
                tracker_dict[matched_id] = {
                    'pos': (cx, cy),
                    't1': None,
                    't2': None,
                    'logged': False,
                    'last_cy': cy,
                    'direction': None
                }
            else:
                tracker_dict[matched_id]['pos'] = (cx, cy)

            vehicle = tracker_dict[matched_id]
            last_cy = vehicle['last_cy']
            vehicle['last_cy'] = cy

            # Downward detection (your original logic)
            if not vehicle['t1'] and cy >= speed_line_y1:
                vehicle['t1'] = datetime.now()
            elif not vehicle['t2'] and cy >= speed_line_y2:
                vehicle['t2'] = datetime.now()
                vehicle['direction'] = 'bawah'

            # Upward detection (added logic)
            if not vehicle['t1'] and cy <= speed_line_y2:
                vehicle['t1'] = datetime.now()
            elif not vehicle['t2'] and cy <= speed_line_y1:
                vehicle['t2'] = datetime.now()
                vehicle['direction'] = 'atas'

            # Speed calculation
            if vehicle['t1'] and vehicle['t2'] and not vehicle['logged']:
                delta_time = (vehicle['t2'] - vehicle['t1']).total_seconds()
                if delta_time > 0:
                    speed_kmh = (meters_travel_pixel / delta_time) * 3.6
                    direction = vehicle['direction'] or 'unknown'
                    print(f"ID {matched_id} SPEED: {speed_kmh:.2f} km/h DIR: {direction}")
                    
                    if speed_kmh >= 10:
                        cursor.execute(
                            "INSERT INTO vehicle_log (vehicle_id, direction, timestamp, speed) VALUES (%s, %s, %s, %s)",
                            (matched_id, direction, datetime.now(), round(speed_kmh, 2))
                        )
                        conn.commit()
                    vehicle['logged'] = True

            # Show ID
            cv2.putText(frame, f"ID:{matched_id}", (x, y - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1)

            # Show speed
            if vehicle['t1'] and vehicle['t2']:
                delta_time = (vehicle['t2'] - vehicle['t1']).total_seconds()
                if delta_time > 0:
                    speed_kmh = (meters_travel_pixel / delta_time) * 3.6
                    cv2.putText(frame, f"{speed_kmh:.1f} km/h", (x, y + h + 15),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)

    
    # Store vehicle count in zone
    if vehicle_ids_in_zone != last_vehicle_ids_in_zone:
        last_vehicle_ids_in_zone = vehicle_ids_in_zone.copy()
        count = len(vehicle_ids_in_zone)
        cursor.execute("INSERT INTO total_vehicle (count, timestamp) VALUES (%s, %s)", (count, datetime.now()))
        conn.commit()
        print(f"🚗 Vehicles in zone: {count}")
    
    # Draw detection lines
    cv2.line(frame, (0, speed_line_y1), (width2, speed_line_y1), (0, 255, 255), 2)
    cv2.line(frame, (0, speed_line_y2), (width2, speed_line_y2), (255, 0, 255), 2)

    cv2.imshow("Vehicle Detection", frame)
    if cv2.waitKey(1) & 0xFF == 27:
        break

cap.release()
cursor.close()
conn.close()
cv2.destroyAllWindows()