import csv
import time
from datetime import datetime
from ultralytics import YOLO
import cv2
import mediapipe as mp
import numpy as np

model = YOLO('best_5.pt')

cap = cv2.VideoCapture(0)

mp_hands = mp.solutions.hands
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

hands = mp_hands.Hands(static_image_mode=True, min_detection_confidence=0.3)
labels_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'H'}

# File to save gesture data
csv_file = 'gesture_data.csv'

# Initialize variables for tracking gestures
previous_gesture = None
gesture_start_time = None

# List to store gesture data for CSV
gesture_data_list = []

while True:
    data_aux = []
    x_ = []
    y_ = []

    ret, frame = cap.read()
    if not ret:
        print("Failed to capture frame. Exiting...")
        break

    H, W, _ = frame.shape

    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    results = hands.process(frame_rgb)
    if results.multi_hand_landmarks:
        for hand_landmarks in results.multi_hand_landmarks:
            mp_drawing.draw_landmarks(
                frame,
                hand_landmarks,
                mp_hands.HAND_CONNECTIONS,
                mp_drawing_styles.get_default_hand_landmarks_style(),
                mp_drawing_styles.get_default_hand_connections_style())

        for hand_landmarks in results.multi_hand_landmarks:
            for i in range(len(hand_landmarks.landmark)):
                x = hand_landmarks.landmark[i].x
                y = hand_landmarks.landmark[i].y

                x_.append(x)
                y_.append(y)

            for i in range(len(hand_landmarks.landmark)):
                x = hand_landmarks.landmark[i].x
                y = hand_landmarks.landmark[i].y
                data_aux.append(x - min(x_))
                data_aux.append(y - min(y_))

        x1 = int(min(x_) * W) - 10
        y1 = int(min(y_) * H) - 10

        x2 = int(max(x_) * W) + 10
        y2 = int(max(y_) * H) + 10

        prediction = model.predict(frame, conf=0.25, iou=0.45)
        names_dict = prediction[0].names
        probs = prediction[0].probs.data.numpy()
        detected_gesture = names_dict[np.argmax(probs)]
        print("Gesture:", detected_gesture)

        if detected_gesture == 'A':
            language = 'Arabic'
        elif detected_gesture == 'B':
            language = 'Bengali'
        elif detected_gesture == 'C':
            language = 'Chinese'
        elif detected_gesture == 'H':
            language = 'Hindi'

        # Get the current timestamp
        current_time = time.time()

        # Check if the detected gesture has changed
        if detected_gesture != previous_gesture:
            # If the detected gesture has changed, calculate the duration of the previous gesture
            if previous_gesture is not None:
                gesture_end_time = current_time
                gesture_duration = gesture_end_time - gesture_start_time
                # Store the detected gesture, start time, end time, and duration in the list
                gesture_data_list.append([previous_gesture,
                                          datetime.fromtimestamp(gesture_start_time).strftime('%H:%M:%S.%f'),
                                          datetime.fromtimestamp(gesture_end_time).strftime('%H:%M:%S.%f'),
                                          round(gesture_duration, 2)])

            # Update the previous gesture and its start time
            previous_gesture = detected_gesture
            gesture_start_time = current_time

        # Calculate the duration of the current gesture
        gesture_duration = current_time - gesture_start_time

        # Draw rectangle around the detected gesture
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 0), 4)
        cv2.putText(frame, detected_gesture, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)

    # Display the frame
    cv2.imshow('frame', frame)

    # Check if 'q' key is pressed to stop the program
    if cv2.waitKey(1) & 0xFF == ord('q'):
        # Save data to CSV file
        with open(csv_file, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['Gesture', 'Start Time', 'End Time', 'Duration'])
            for gesture, start_time, end_time, duration in gesture_data_list:
                writer.writerow([gesture, start_time, end_time, duration])
        break

# Release resources
cap.release()
cv2.destroyAllWindows()