import cv2
import mediapipe as mp
import csv
import numpy as np

# Define a 1D Kalman Filter class with velocity
class KalmanFilter1D:
    def __init__(self, initial_position=0.0):
        # State vector [position, velocity]
        self.state = np.array([[initial_position],
                               [0.0]], dtype=np.float32)
        # State covariance matrix
        self.covariance = np.eye(2, dtype=np.float32)
        # State transition matrix (assuming constant velocity)
        self.F = np.array([[1, 1],
                           [0, 1]], dtype=np.float32)
        # Measurement matrix (we only measure position)
        self.H = np.array([[1, 0]], dtype=np.float32)
        # Process noise covariance
        self.Q = np.array([[1e-4, 0],
                           [0, 1e-4]], dtype=np.float32)
        # Measurement noise covariance
        self.R = np.array([[1e-1]], dtype=np.float32)
        # Initialization flag
        self.initialized = False

    def initialize(self, measurement):
        self.state = np.array([[measurement],
                               [0.0]], dtype=np.float32)
        self.initialized = True

    def predict(self):
        # Predict the next state
        self.state = self.F @ self.state
        # Update the state covariance matrix
        self.covariance = self.F @ self.covariance @ self.F.T + self.Q

    def update(self, measurement):
        if not self.initialized:
            self.initialize(measurement)
            return self.state[0][0], self.state[1][0]  # Position, Velocity

        # Measurement residual
        y = np.array([[measurement]]) - (self.H @ self.state)
        # Measurement covariance
        S = self.H @ self.covariance @ self.H.T + self.R
        # Kalman gain
        K = self.covariance @ self.H.T @ np.linalg.inv(S)
        # Update the state estimate
        self.state = self.state + K @ y
        # Update the covariance estimate
        I = np.eye(self.F.shape[0], dtype=np.float32)
        self.covariance = (I - K @ self.H) @ self.covariance

        return self.state[0][0], self.state[1][0]  # Position, Velocity

# Initialize MediaPipe Hand Tracking
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(
    static_image_mode=False,       # For video input
    max_num_hands=2,               # Maximum number of hands to detect
    min_detection_confidence=0.5,  # Minimum confidence for detection
    min_tracking_confidence=0.5    # Minimum confidence for tracking
)
mp_drawing = mp.solutions.drawing_utils  # For drawing hand landmarks

# Open the input video
cap = cv2.VideoCapture('Test_video.mp4')

# Get video properties
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))   # Video frame width
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # Video frame height
fps = cap.get(cv2.CAP_PROP_FPS)                  # Original frames per second

# Desired processing frame rate
desired_fps = 10
frame_skip = max(int(fps / desired_fps), 1)  # Ensure at least 1
delta_t = 1 / desired_fps

# Initialize Video Writer for output video
output_video = cv2.VideoWriter(
    'output_video_kf.mp4',
    cv2.VideoWriter_fourcc(*'mp4v'),
    desired_fps,  # Set output video FPS to desired FPS
    (width, height)
)

# Create CSV header
header = ['Frame', 'Hand']
for i in range(21):  # 21 landmarks in MediaPipe Hands
    header.extend([
        f'Landmark_{i}_x', f'Landmark_{i}_y', f'Landmark_{i}_z',
        f'Landmark_{i}_velocity_x', f'Landmark_{i}_velocity_y', f'Landmark_{i}_velocity_z',
        f'Landmark_{i}_kalman_x', f'Landmark_{i}_kalman_y', f'Landmark_{i}_kalman_z'
    ])

# Initialize Kalman filters for all 21 landmarks (for each hand)
def initialize_kalman_filters():
    return {
        'Left': [{'x': KalmanFilter1D(), 'y': KalmanFilter1D(), 'z': KalmanFilter1D()} for _ in range(21)],
        'Right': [{'x': KalmanFilter1D(), 'y': KalmanFilter1D(), 'z': KalmanFilter1D()} for _ in range(21)]
    }

kalman_filters = initialize_kalman_filters()

# Store previous landmarks to compute velocity
previous_landmarks = {'Left': None, 'Right': None}

# Initialize frames since detection counters for left and right hands
frames_since_detection = {'Left': 0, 'Right': 0}

# Maximum allowed frames without detection before resetting filters
max_frames_without_detection = 10

frame_count = 0  # To keep track of the frame number

def process_hand(hand_label, hand_landmarks, previous_landmarks, kalman_filters, frame_count, annotated_frame):
    # Get the wrist landmark as reference
    wrist = hand_landmarks.landmark[mp_hands.HandLandmark.WRIST]
    wrist_x = wrist.x
    wrist_y = wrist.y
    wrist_z = wrist.z

    # Initialize a list to store adjusted landmarks
    landmarks_row = [frame_count, hand_label]
    current_landmarks = []  # To store current frame's landmarks

    for i, landmark in enumerate(hand_landmarks.landmark):
        # Original normalized coordinates
        x_raw = landmark.x
        y_raw = landmark.y
        z_raw = landmark.z

        # Adjust landmarks relative to wrist
        x = x_raw - wrist_x
        y = y_raw - wrist_y
        z = z_raw - wrist_z

        # Calculate velocity using previous landmarks
        if previous_landmarks[hand_label] is not None:
            previous_x, previous_y, previous_z = previous_landmarks[hand_label][i]
            velocity_x = (x - previous_x) / delta_t
            velocity_y = (y - previous_y) / delta_t
            velocity_z = (z - previous_z) / delta_t
        else:
            velocity_x = 0.0
            velocity_y = 0.0
            velocity_z = 0.0

        current_landmarks.append((x, y, z))

        # Apply Kalman filter to the x, y, and z coordinates
        kalman_x = kalman_filters[hand_label][i]['x']
        kalman_y = kalman_filters[hand_label][i]['y']
        kalman_z = kalman_filters[hand_label][i]['z']

        if not kalman_x.initialized:
            # Initialize filters with the current measurement
            kalman_x.initialize(x)
            kalman_y.initialize(y)
            kalman_z.initialize(z)

        # Prediction step
        kalman_x.predict()
        kalman_y.predict()
        kalman_z.predict()

        # Update step with measurement
        kalman_x_estimated, _ = kalman_x.update(x)
        kalman_y_estimated, _ = kalman_y.update(y)
        kalman_z_estimated, _ = kalman_z.update(z)

        # Visualization of Raw and Kalman-Filtered Landmarks
        # Convert normalized coordinates to pixel coordinates for visualization
        raw_x_pixel = int(x_raw * width)
        raw_y_pixel = int(y_raw * height)
        kalman_x_abs = kalman_x_estimated + wrist_x
        kalman_y_abs = kalman_y_estimated + wrist_y
        kalman_x_pixel = int(kalman_x_abs * width)
        kalman_y_pixel = int(kalman_y_abs * height)

        # Ensure pixel values are within image bounds
        raw_x_pixel = np.clip(raw_x_pixel, 0, width - 1)
        raw_y_pixel = np.clip(raw_y_pixel, 0, height - 1)
        kalman_x_pixel = np.clip(kalman_x_pixel, 0, width - 1)
        kalman_y_pixel = np.clip(kalman_y_pixel, 0, height - 1)

        # Draw raw landmark (green)
        cv2.circle(annotated_frame, (raw_x_pixel, raw_y_pixel), 3, (0, 255, 0), -1)
        # Draw Kalman-filtered landmark (red)
        cv2.circle(annotated_frame, (kalman_x_pixel, kalman_y_pixel), 3, (0, 0, 255), -1)

        # Append data to CSV row
        landmarks_row.extend([
            x, y, z,
            velocity_x, velocity_y, velocity_z,
            kalman_x_estimated, kalman_y_estimated, kalman_z_estimated
        ])

    # Update previous landmarks
    previous_landmarks[hand_label] = current_landmarks
    return landmarks_row

def predict_missing_hand(hand_label, kalman_filters, frames_since_detection, frame_count):
    # Only predict for a certain number of frames after the hand is missing
    if frames_since_detection[hand_label] < max_frames_without_detection:
        landmarks_row = [frame_count, hand_label]
        for i in range(21):
            kalman_x = kalman_filters[hand_label][i]['x']
            kalman_y = kalman_filters[hand_label][i]['y']
            kalman_z = kalman_filters[hand_label][i]['z']

            if kalman_x.initialized:
                # Prediction step
                kalman_x.predict()
                kalman_y.predict()
                kalman_z.predict()

                # Use predicted state as current estimate (no measurement update)
                kalman_x_estimated = kalman_x.state[0][0]
                kalman_y_estimated = kalman_y.state[0][0]
                kalman_z_estimated = kalman_z.state[0][0]

                # Append empty measurements and velocities, but include Kalman estimates
                landmarks_row.extend([
                    '', '', '',
                    '', '', '',
                    kalman_x_estimated, kalman_y_estimated, kalman_z_estimated
                ])
            else:
                # Append empty data if Kalman filter is not initialized
                landmarks_row.extend([''] * 9)

        frames_since_detection[hand_label] += 1
        # Reset Kalman filters if the hand is missing for too long
        if frames_since_detection[hand_label] >= max_frames_without_detection:
            for i in range(21):
                kalman_filters[hand_label][i]['x'] = KalmanFilter1D()
                kalman_filters[hand_label][i]['y'] = KalmanFilter1D()
                kalman_filters[hand_label][i]['z'] = KalmanFilter1D()
            previous_landmarks[hand_label] = None
            frames_since_detection[hand_label] = 0
        return landmarks_row
    return None

# Prepare to write to CSV file
with open('landmarks_with_velocity_and_kalman.csv', mode='w', newline='') as csv_file:
    csv_writer = csv.writer(csv_file)
    csv_writer.writerow(header)

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

        if frame_count % frame_skip == 0:
            # Convert frame to RGB for MediaPipe processing
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            result = hands.process(frame_rgb)

            # Create a copy of the frame for drawing
            annotated_frame = frame.copy()
            csv_rows = []

            if result.multi_hand_landmarks and result.multi_handedness:
                detected_hands = []
                for hand_landmarks, handedness in zip(result.multi_hand_landmarks, result.multi_handedness):
                    hand_label = handedness.classification[0].label
                    detected_hands.append(hand_label)

                    # Draw the hand landmarks on the frame
                    mp_drawing.draw_landmarks(
                        annotated_frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

                    # Reset frames since detection for this hand
                    frames_since_detection[hand_label] = 0

                    # Process hand and collect data
                    landmarks_row = process_hand(hand_label, hand_landmarks, previous_landmarks, kalman_filters, frame_count, annotated_frame)
                    csv_rows.append(landmarks_row)

                # Handle hands not detected in this frame
                for hand_label in ['Left', 'Right']:
                    if hand_label not in detected_hands:
                        missing_hand_row = predict_missing_hand(hand_label, kalman_filters, frames_since_detection, frame_count)
                        if missing_hand_row:
                            csv_rows.append(missing_hand_row)
            else:
                # No hands detected in this frame
                for hand_label in ['Left', 'Right']:
                    missing_hand_row = predict_missing_hand(hand_label, kalman_filters, frames_since_detection, frame_count)
                    if missing_hand_row:
                        csv_rows.append(missing_hand_row)

            # Write all rows for this frame to the CSV
            for row in csv_rows:
                csv_writer.writerow(row)

            # Write the annotated frame to output video
            output_video.write(annotated_frame)

        frame_count += 1  # Increment frame count regardless of processing

# Release resources
cap.release()
output_video.release()
hands.close()
add the following below features for detect and track and save landmarks on black gloves frames: 
Primary Detection:

Use MOG2 for background subtraction due to its adaptability to dynamic backgrounds and varying lighting conditions. It can handle the challenges of a moving camera better than KNN.
Enhancement & Cleanup:

Apply morphological operations (erosion and dilation) to clean up noise from the MOG2 mask. This will ensure that small false positives or noise are removed, making the gloves' contours clearer.
Color Filtering:

Use HSV color filtering to isolate the black gloves. This will enhance detection when the gloves are distinguishable by their color in the HSV space (for example, matte black gloves).
Edge Detection (Optional):

Edge detection can be added to further refine the boundaries of the gloves, ensuring that they are well-outlined, especially if the gloves have distinct edges against the background.
Tracking:

Kalman filtering should be applied to smooth the glove trajectory, especially in cases where detection might falter due to camera shake or momentary occlusion. This will help maintain consistent tracking over time.(use mediapipe instead of kalman filter if better orther wise use kalman filter)?
use this approach or do any changes as necessary for better prediction of black gloves?