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')  # Change to RTSP URL if needed

# Resize ratio
ratio = 0.5
ret, frame = cap.read()
image = cv2.resize(frame, (0, 0), None, ratio, ratio)
width2, height2, channels = image.shape

# Background subtractor
fgbg = cv2.createBackgroundSubtractorMOG2()

# Tracking + speed
vehicle_id_counter = 0
tracker_dict = {}
speed_line_y1 = 150
speed_line_y2 = 200
pixel_distance = abs(speed_line_y2 - speed_line_y1)
fps = cap.get(cv2.CAP_PROP_FPS) or 60  # fallback default

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

    image = cv2.resize(frame, (0, 0), None, ratio, ratio)
    gray = cv2.cvtColor(image, 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(image, hull, -1, (0, 255, 0), 2)

    min_area = 400
    max_area = 40000

    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'])

            x, y, w, h = cv2.boundingRect(cnt)
            cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)
            cv2.drawMarker(image, (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}
            else:
                tracker_dict[matched_id]['pos'] = (cx, cy)

            vehicle = tracker_dict[matched_id]
            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()

            if vehicle['t1'] and vehicle['t2'] and not vehicle['logged']:
                delta_time = (vehicle['t2'] - vehicle['t1']).total_seconds()
                if delta_time > 0:
                    meters = 1 # realistic estimated distance in meters
                    speed_kmh = (meters / delta_time) * 3.6

                    print(f"ID {matched_id} SPEED: {speed_kmh:.2f} km/h")

                    # Log to database
                    cursor.execute(
                        "INSERT INTO vehicle_log (vehicle_id, direction, timestamp, speed) VALUES (%s, %s, %s, %s)",
                        (matched_id, 'down', datetime.now(), round(speed_kmh, 2))
                    )
                    conn.commit()
                    vehicle['logged'] = True

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

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

    # Draw full-width lines
    cv2.line(image, (0, speed_line_y1), (width2, speed_line_y1), (0, 255, 255), 2)
    cv2.line(image, (0, speed_line_y2), (width2, speed_line_y2), (255, 0, 255), 2)

    # Display window
    cv2.imshow("Vehicle Detection", image)
    if cv2.waitKey(1) & 0xFF == 27:  # ESC to quit
        break

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