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

# Add src to path
sys.path.append(str(Path(__file__).parent.parent.parent))
from src.utils.preprocess import GazePreprocessor

def process_gaze360_subset(mat_path, img_root, output_path, rec_idx=6, patch_size=16):
    print(f"Loading metadata from {mat_path}...")
    mat = scipy.io.loadmat(mat_path)
    
    recording_indices = mat['recording'].flatten()
    mask = (recording_indices == rec_idx)
    
    indices = np.where(mask)[0]
    print(f"Found {len(indices)} potential samples in recording {rec_idx}")
    
    preprocessor = GazePreprocessor()
    
    # Get recording name
    rec_name = mat['recordings'][0, rec_idx][0]
    
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    with h5py.File(output_path, 'w') as h5f:
        # Datasets
        left_patches_ds = h5f.create_dataset('left_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        right_patches_ds = h5f.create_dataset('right_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        # We store person's gaze as both left and right for compatibility
        gaze_ds = h5f.create_dataset('gaze', (0, 2), maxshape=(None, 2), dtype='float32') 
        landmarks_ds = h5f.create_dataset('landmarks', (0, 478, 2), maxshape=(None, 478, 2), dtype='float32')
        
        sample_idx = 0
        
        for i in tqdm(indices, desc="Processing Gaze360"):
            person_id = mat['person_identity'][0, i]
            frame_num = mat['frame'][0, i]
            gaze_3d = mat['gaze_dir'][i] # (3,)
            
            # Construct path
            img_path = os.path.join(img_root, rec_name, 'head', f'{person_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(frame)
            if landmarks is None:
                continue
            
            # 1. Landmarks (Centered)
            landmarks_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 = landmarks_arr - face_center
            
            # 2. Eye Patches
            try:
                # We use Left Eye for normalizing the 3D gaze (assuming they are close)
                # In a more precise setup, we'd do it for each eye.
                left_eye_img, left_angle = preprocessor.normalize_eye(frame, landmarks, 'left')
                left_patches = preprocessor.extract_patches(left_eye_img, patch_size=patch_size)
                
                right_eye_img, right_angle = preprocessor.normalize_eye(frame, landmarks, 'right')
                right_patches = preprocessor.extract_patches(right_eye_img, patch_size=patch_size)
                
                # Rotate Gaze (Using average angle or just left eye's angle)
                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) # (pitch, yaw)
                
                # Store
                for ds, data in zip([left_patches_ds, right_patches_ds, gaze_ds, landmarks_ds],
                                   [left_patches, right_patches, gaze_rad, landmarks_centered]):
                    ds.resize((sample_idx + 1, *ds.shape[1:]))
                    ds[sample_idx] = data
                
                sample_idx += 1
            except Exception as e:
                # print(f"Error processing frame: {e}")
                continue
                
    print(f"\nFinished! Processed {sample_idx} samples. Saved to {output_path}")

if __name__ == "__main__":
    process_gaze360_subset(
        mat_path='data/raw/metadata.mat',
        img_root='data/raw/imgs',
        output_path='data/processed/gaze360_v16.h5',
        rec_idx=6,
        patch_size=16
    )