import os
import cv2
import numpy as np
import h5py
import scipy.io
from tqdm import tqdm
import sys
from pathlib import Path

# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
from src.utils.preprocess import GazePreprocessor

class RobustGazePreprocessor(GazePreprocessor):
    def get_landmarks_robust(self, frame):
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        landmarks = self.get_landmarks(frame)
        if landmarks: return landmarks
        return None

def create_unseen_split():
    mat_path = 'data/raw/metadata.mat'
    img_root = 'data/raw/imgs'
    output_path = 'data/processed/gaze360_unseen_verification.h5'
    
    # Target IDs in Recording 6 that were NOT in test_B
    # rec_006 IDs: [0, 1, 17, 25, 49, 60, 61, 62]
    # test_B had: [60, 62]
    UNSEEN_IDS = [0, 1, 17, 25, 49, 61]
    REC_IDX = 6
    patch_size = 16

    print(f"Loading metadata and selecting unseen IDs: {UNSEEN_IDS}")
    mat = scipy.io.loadmat(mat_path)
    rec_indices = mat['recording'].flatten()
    person_ids = mat['person_identity'].flatten()
    
    mask = (rec_indices == REC_IDX) & np.isin(person_ids, UNSEEN_IDS)
    indices = np.where(mask)[0]
    rec_name = mat['recordings'][0, REC_IDX][0]
    
    preprocessor = RobustGazePreprocessor()
    os.makedirs('data/processed', exist_ok=True)

    with h5py.File(output_path, 'w') as h5f:
        lp_ds = h5f.create_dataset('left_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        rp_ds = h5f.create_dataset('right_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        gaze_ds = h5f.create_dataset('gaze', (0, 2), maxshape=(None, 2), dtype='float32')
        lm_ds = h5f.create_dataset('landmarks', (0, 478, 2), maxshape=(None, 478, 2), dtype='float32')
        
        count = 0
        for i in tqdm(indices, desc="Preprocessing Unseen Subjects"):
            p_id = person_ids[i]
            frame_num = mat['frame'][0, i]
            gaze_3d = mat['gaze_dir'][i]
            
            img_path = os.path.join(img_root, rec_name, 'head', f'{p_id:06d}', f'{frame_num:06d}.jpg')
            if not os.path.exists(img_path): continue
                
            frame = cv2.imread(img_path)
            if frame is None: continue
                
            landmarks = preprocessor.get_landmarks_robust(frame)
            if landmarks is None: continue

            try:
                # Normalize and extract patches
                left_eye_img, left_angle = preprocessor.normalize_eye(frame, landmarks, 'left', method='new')
                left_patches = preprocessor.extract_patches(left_eye_img, patch_size=patch_size)
                
                right_eye_img, right_angle = preprocessor.normalize_eye(frame, landmarks, 'right', method='new')
                right_patches = preprocessor.extract_patches(right_eye_img, patch_size=patch_size)
                
                # Align Gaze
                avg_angle = (left_angle + right_angle) / 2
                gaze_rot = preprocessor.rotate_gaze(gaze_3d, avg_angle)
                gaze_rad = preprocessor.gaze_3d_to_mag(gaze_rot)

                # Landmarks relative to face center
                lms_arr = np.array([[lm.x, lm.y] for lm in landmarks])
                left_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in preprocessor.LEFT_CORNERS], axis=0)
                right_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in preprocessor.RIGHT_CORNERS], axis=0)
                face_center = (left_c + right_c) / 2
                landmarks_centered = lms_arr - face_center

                # Save
                lp_ds.resize((count + 1, 4, patch_size, patch_size))
                lp_ds[count] = left_patches
                rp_ds.resize((count + 1, 4, patch_size, patch_size))
                rp_ds[count] = right_patches
                gaze_ds.resize((count + 1, 2))
                gaze_ds[count] = gaze_rad
                lm_ds.resize((count + 1, 478, 2))
                lm_ds[count] = landmarks_centered
                
                count += 1
            except:
                continue

    print(f"\nDone! Created {output_path} with {count} samples.")

if __name__ == "__main__":
    create_unseen_split()