import cv2
import mediapipe as mp
import numpy as np
import pandas as pd
import os
from sklearn.cluster import KMeans

class PersistentSignSegmenter:
    def __init__(self, start_threshold=0.015, stop_threshold=0.01, temporal_padding=15):
        self.mp_holistic = mp.solutions.holistic
        self.start_threshold = start_threshold
        self.stop_threshold = stop_threshold
        self.temporal_padding = temporal_padding
        self.emotions = ["neutral", "happy", "sad", "anger", "disgust", "fear", "surprise"]
        self.target_size = (640, 480)
        self.log_file = "processed_videos.log"
        self.master_csv = "master_annotations.csv"

    def _resize_with_pad(self, image):
        h, w = image.shape[:2]
        target_w, target_h = self.target_size
        scale = min(target_w / w, target_h / h)
        nw, nh = int(w * scale), int(h * scale)
        
        image_resized = cv2.resize(image, (nw, nh))
        canvas = np.zeros((target_h, target_w, 3), dtype=np.uint8)
        
        offset_x, offset_y = (target_w - nw) // 2, (target_h - nh) // 2
        canvas[offset_y:offset_y+nh, offset_x:offset_x+nw] = image_resized
        return canvas

    def _get_processed_list(self):
        if not os.path.exists(self.log_file):
            return set()
        with open(self.log_file, "r") as f:
            return set(line.strip() for line in f)

    def _log_progress(self, filename):
        with open(self.log_file, "a") as f:
            f.write(f"{filename}\n")

    def _calculate_velocity(self, results, prev_coords):
        vel = 0
        active_hands = 0
        for landmarks in [results.left_hand_landmarks, results.right_hand_landmarks]:
            if landmarks:
                coords = np.array([[lm.x, lm.y] for lm in landmarks.landmark])
                if prev_coords is not None and prev_coords.shape == coords.shape:
                    dist = np.mean(np.linalg.norm(coords - prev_coords, axis=1))
                    vel += dist
                prev_coords = coords
                active_hands += 1
        return vel if active_hands > 0 else 0, prev_coords

    def process_folder(self, input_folder):
        processed_files = self._get_processed_list()
        video_files = [f for f in os.listdir(input_folder) if f.endswith(('.mp4', '.avi', '.mov'))]
        
        for video_file in video_files:
            if video_file in processed_files:
                print(f"[SKIP] {video_file} already processed.")
                continue

            video_path = os.path.join(input_folder, video_file)
            print(f"[*] Processing: {video_file}")
            
            segments = self._get_segments(video_path)
            
            if len(segments) >= 7:
                midpoints = np.array([(s['start_frame'] + s['end_frame']) / 2 for s in segments]).reshape(-1, 1)
                kmeans = KMeans(n_clusters=7, random_state=42, n_init=10).fit(midpoints)
                cluster_map = np.argsort(kmeans.cluster_centers_.flatten())
                label_mapping = {cluster_idx: self.emotions[i] for i, cluster_idx in enumerate(cluster_map)}
                
                new_rows = []
                for i, seg in enumerate(segments):
                    seg['label'] = label_mapping[kmeans.labels_[i]]
                    seg['filename'] = video_file
                    new_rows.append(seg)
                
                df = pd.DataFrame(new_rows)
                df.to_csv(self.master_csv, mode='a', header=not os.path.exists(self.master_csv), index=False)
                
                self._log_progress(video_file)
                print(f"[DONE] {video_file} - {len(segments)} segments identified.")
            else:
                print(f"[ERROR] {video_file} only had {len(segments)} segments. Check thresholds.")

    def _get_segments(self, video_path):
        cap = cv2.VideoCapture(video_path)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        segments = []
        is_recording, current_start, quiet_counter, prev_coords = False, 0, 0, None

        with self.mp_holistic.Holistic(min_detection_confidence=0.5) as holistic:
            for i in range(total_frames):
                ret, frame = cap.read()
                if not ret: break
                
                # Resize and Pad
                processed_frame = self._resize_with_pad(frame)
                
                res = holistic.process(cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB))
                vel, prev_coords = self._calculate_velocity(res, prev_coords)
                
                if not is_recording and vel > self.start_threshold:
                    is_recording, current_start = True, i
                elif is_recording:
                    if vel < self.stop_threshold: quiet_counter += 1
                    else: quiet_counter = 0
                    
                    if quiet_counter > 15:
                        segments.append({
                            'start_frame': max(0, current_start - self.temporal_padding),
                            'end_frame': min(total_frames - 1, (i - 15) + self.temporal_padding)
                        })
                        is_recording = False
                        quiet_counter = 0
        cap.release()
        return segments

processor = PersistentSignSegmenter()
processor.process_folder("fsl-data/raw_videos")