import warnings
warnings.filterwarnings("ignore", category=UserWarning)

import cv2
import csv
import os
import numpy as np
from picamera2 import Picamera2
import mediapipe.python.solutions.hands as mp_hands
import mediapipe.python.solutions.drawing_utils as mp_draw

LABELS = {
    ord('1'): "FIST",
    ord('2'): "OPEN_PALM",
    ord('3'): "POINTING",
    ord('4'): "THREE_FINGERS",
    ord('5'): "PEACE",
    ord('6'): "FOUR_FINGERS",
    ord('0'): "NONE",
}

CSV_FILE = "gesture_data.csv"

hands = mp_hands.Hands(max_num_hands=1, model_complexity=0,
    min_detection_confidence=0.6, min_tracking_confidence=0.5)

def normalize(lm):
    pts = np.array([[p.x, p.y, p.z] for p in lm])
    pts = pts - pts[0]
    scale = np.linalg.norm(pts[9])
    pts = pts / (scale + 1e-6)
    return pts.flatten()

if not os.path.exists(CSV_FILE):
    with open(CSV_FILE, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow([f"f{i}" for i in range(63)] + ["label"])

picam2 = Picamera2()
picam2.configure(picam2.create_preview_configuration(
    main={"format": "RGB888", "size": (640, 480)}))
picam2.start()

recording_label = None
counts = {}

print("Keys: 1=FIST 2=OPEN_PALM 3=POINTING 4=THREE_FINGERS 5=PEACE 6=FOUR_FINGERS 0=NONE")
print("s = stop recording, q = quit")

while True:
    frame = picam2.capture_array()          # RGB
    results = hands.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    display = frame.copy()

    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):
        break
    if key == ord('s'):
        recording_label = None
    if key in LABELS:
        recording_label = LABELS[key]

    if results.multi_hand_landmarks:
        h = results.multi_hand_landmarks[0]
        mp_draw.draw_landmarks(display, h, mp_hands.HAND_CONNECTIONS)
        if recording_label is not None:
            feat = normalize(h.landmark)
            with open(CSV_FILE, "a", newline="") as f:
                csv.writer(f).writerow(list(feat) + [recording_label])
            counts[recording_label] = counts.get(recording_label, 0) + 1

    status = "NOT RECORDING (press 1-5 or 0)"
    color = (0, 0, 255)
    if recording_label:
        n = counts.get(recording_label, 0)
        status = f"RECORDING {recording_label}: {n} samples (s=stop)"
        color = (0, 255, 0)
    cv2.putText(display, status, (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

    cv2.imshow("Data Collection (Pi)", display)

picam2.stop()
cv2.destroyAllWindows()
print("Saved counts:", counts)