import cv2
import time
import os
import threading
from datetime import datetime
from feat import Detector
import torch
import mediapipe as mp
import csv

# Konfiguration
CAMERA_INDEX = 0
OUTPUT_DIR = "recordings"
VIDEO_DURATION = 10  # Sekunden
START_INTERVAL = 5   # Sekunden bis zum nächsten Start
FPS = 25.0           # Feste FPS

if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

# Globaler Detector, um ihn nicht bei jedem Video neu laden zu müssen (spart massiv Zeit/Speicher)
print("Initialisiere AU-Detector (bitte warten)...")
detector = Detector(au_model="xgb") 

# ===== MediaPipe FaceMesh Setup =====
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(
    static_image_mode=False,
    max_num_faces=1,
    refine_landmarks=True,  # wichtig für Iris
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5
)

LEFT_IRIS = [474, 475, 476, 477]
RIGHT_IRIS = [469, 470, 471, 472]

LEFT_EYE_LIDS = (159, 145)
RIGHT_EYE_LIDS = (386, 374)

LEFT_EYE_GAZE_IDXS = (33, 133, 159, 145)
RIGHT_EYE_GAZE_IDXS = (263, 362, 386, 374)

EYE_OPEN_THRESHOLD = 6

# CSV vorbereiten
gaze_csv = open("gaze_data.csv", mode="w", newline="")
gaze_writer = csv.writer(gaze_csv)
gaze_writer.writerow([
    "timestamp",
    "left_gaze_x",
    "left_gaze_y",
    "right_gaze_x",
    "right_gaze_y",
    "left_valid",
    "right_valid",
    "left_diameter",
    "right_diameter"
])

def eye_openness(landmarks, top_idx, bottom_idx, img_height):
    top = landmarks[top_idx]
    bottom = landmarks[bottom_idx]
    return abs(top.y - bottom.y) * img_height


def compute_gaze(landmarks, iris_center, indices, w, h):
    idx1, idx2, top_idx, bottom_idx = indices

    p1 = landmarks[idx1]
    p2 = landmarks[idx2]
    top = landmarks[top_idx]
    bottom = landmarks[bottom_idx]

    x1 = p1.x * w
    x2 = p2.x * w
    y_top = top.y * h
    y_bottom = bottom.y * h

    iris_x, iris_y = iris_center

    eye_left = min(x1, x2)
    eye_right = max(x1, x2)

    eye_width = eye_right - eye_left
    eye_height = abs(y_bottom - y_top)

    if eye_width == 0 or eye_height == 0:
        return 0.5, 0.5

    gaze_x = (iris_x - eye_left) / eye_width
    gaze_y = (iris_y - min(y_top, y_bottom)) / eye_height

    gaze_x = max(0, min(1, gaze_x))
    gaze_y = max(0, min(1, gaze_y))

    return gaze_x, gaze_y

def extract_aus(path, skip_frames):

    # torch.no_grad() deaktiviert die Gradientenberechnung.
    # Das löst den "Can't call numpy() on Tensor that requires grad" Fehler.
    with torch.no_grad():
        video_prediction = detector.detect_video(
            path, 
            skip_frames=skip_frames, 
            face_detection_threshold=0.95
        )
        
        # Falls video_prediction oder .aus noch Tensoren sind, 
        # stellen wir sicher, dass sie korrekt summiert werden.
        try:
            # Wir nehmen die Summe der Action Units über alle detektierten Frames
            res = video_prediction.aus.sum()
            return res
        except Exception as e:
            print(f"Fehler bei der Summenbildung: {e}")
            return 0

def startAU_creation(video_path):
    """Diese Funktion läuft nun in einem eigenen Thread."""
    try:
        print(f"\n[THREAD START] Analyse läuft für: {video_path}")
        # skip_frames berechnen (z.B. alle 5 Sekunden bei 25 FPS = 125)
        output = extract_aus(video_path, skip_frames=int(FPS*5))
        
        print(f"\n--- Ergebnis für {os.path.basename(video_path)} ---")
        print(output)
        print("--------------------------------------------------\n")
    except Exception as e:
        print(f"Fehler bei der Analyse von {video_path}: {e}")

class VideoRecorder:
    def __init__(self, filename, width, height):
        self.filename = filename
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.out = cv2.VideoWriter(filename, fourcc, FPS, (width, height))
        self.frames_to_record = int(VIDEO_DURATION * FPS)
        self.frames_count = 0
        self.is_finished = False

    def write_frame(self, frame):
        if self.frames_count < self.frames_to_record:
            self.out.write(frame)
            self.frames_count += 1
        else:
            self.finish()

    def finish(self):
        if not self.is_finished:
            self.out.release()
            self.is_finished = True
            abs_path = os.path.abspath(self.filename)
            print(f"Video fertig gespeichert: {self.filename}")

            # --- MULTITHREADING HIER ---
            # Wir starten die Analyse in einem neuen Thread, damit main() sofort weiter frames lesen kann
            analysis_thread = threading.Thread(target=startAU_creation, args=(abs_path,))
            analysis_thread.daemon = True  # Beendet sich, wenn das Hauptprogramm schließt
            analysis_thread.start()

def main():
    cap = cv2.VideoCapture(CAMERA_INDEX)
    if not cap.isOpened():
        print("Fehler: Kamera konnte nicht geöffnet werden.")
        return

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    active_recorders = []
    last_start_time = 0
    
    print("Aufnahme läuft. Drücke 'q' zum Beenden.")

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

            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, _ = frame.shape
            results = face_mesh.process(rgb)

            left_valid = 0
            right_valid = 0
            left_diameter = None
            right_diameter = None

            left_gaze_x = None
            left_gaze_y = None
            right_gaze_x = None
            right_gaze_y = None

            if results.multi_face_landmarks:
                face_landmarks = results.multi_face_landmarks[0]

                left_open = eye_openness(
                    face_landmarks.landmark,
                    LEFT_EYE_LIDS[0],
                    LEFT_EYE_LIDS[1],
                    h
                )

                right_open = eye_openness(
                    face_landmarks.landmark,
                    RIGHT_EYE_LIDS[0],
                    RIGHT_EYE_LIDS[1],
                    h
                )

                left_valid = 1 if left_open > EYE_OPEN_THRESHOLD else 0
                right_valid = 1 if right_open > EYE_OPEN_THRESHOLD else 0

                for eye_name, eye_indices in [("left", LEFT_IRIS), ("right", RIGHT_IRIS)]:
                    iris_points = []

                    for idx in eye_indices:
                        lm = face_landmarks.landmark[idx]
                        x_i, y_i = int(lm.x * w), int(lm.y * h)
                        iris_points.append((x_i, y_i))

                    if len(iris_points) == 4:
                        cx = int(sum(p[0] for p in iris_points) / 4)
                        cy = int(sum(p[1] for p in iris_points) / 4)

                        radius = max(
                            ((x - cx) ** 2 + (y - cy) ** 2) ** 0.5
                            for (x, y) in iris_points
                        )

                        diameter = 2 * radius

                        cv2.circle(frame, (cx, cy), int(radius), (0, 255, 0), 2)

                        if eye_name == "left" and left_valid:
                            left_diameter = diameter
                            left_gaze_x, left_gaze_y = compute_gaze(
                                face_landmarks.landmark,
                                (cx, cy),
                                LEFT_EYE_GAZE_IDXS,
                                w, h
                            )

                        elif eye_name == "right" and right_valid:
                            right_diameter = diameter
                            right_gaze_x, right_gaze_y = compute_gaze(
                                face_landmarks.landmark,
                                (cx, cy),
                                RIGHT_EYE_GAZE_IDXS,
                                w, h
                            )

            # CSV schreiben
            gaze_writer.writerow([
                time.time(),
                left_gaze_x,
                left_gaze_y,
                right_gaze_x,
                right_gaze_y,
                left_valid,
                right_valid,
                left_diameter,
                right_diameter
            ])

            current_time = time.time()

            if current_time - last_start_time >= START_INTERVAL:
                timestamp = datetime.now().strftime("%H%M%S")
                filename = os.path.join(OUTPUT_DIR, f"rec_{timestamp}.avi")
                new_recorder = VideoRecorder(filename, width, height)
                active_recorders.append(new_recorder)
                last_start_time = current_time

            for rec in active_recorders[:]:
                rec.write_frame(frame)
                if rec.is_finished:
                    active_recorders.remove(rec)

            cv2.imshow('Kamera Livestream', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

            time.sleep(1/FPS)

    finally:
        gaze_csv.close()
        face_mesh.close()
        cap.release()
        cv2.destroyAllWindows()
        print("Programm beendet. Warte ggf. auf laufende Analysen...")

if __name__ == "__main__":
    main()