import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
import os

class GazePreprocessor:
    def __init__(self, model_path='face_landmarker.task'):
        """
        Initialize MediaPipe Face Landmarker using the Tasks API.
        """
        base_options = python.BaseOptions(model_asset_path=model_path)
        options = vision.FaceLandmarkerOptions(
            base_options=base_options,
            output_face_blendshapes=False,
            output_facial_transformation_matrixes=True,
            num_faces=1
        )
        self.detector = vision.FaceLandmarker.create_from_options(options)
        
        # Huy's Refined CLAHE: Smaller tileGridSize for finer enhancement
        self.clahe = cv2.createCLAHE(clipLimit=1.2, tileGridSize=(4, 4))
        
        # Landmark indices for eyes (MediaPipe Face Mesh / Landmarker)
        self.LEFT_EYE_IDX = [362, 382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386, 385, 384, 398]
        self.RIGHT_EYE_IDX = [33, 7, 163, 144, 145, 153, 154, 155, 133, 173, 157, 158, 159, 160, 161, 246]
        
        # Standardized: p1 = left corner (viewer's left), p2 = right corner (viewer's right)
        # Right Eye (on viewer's left): 33 (outer), 133 (inner)
        # Left Eye (on viewer's right): 362 (inner), 263 (outer)
        self.RIGHT_CORNERS = [33, 133] 
        self.LEFT_CORNERS = [362, 263] 

    def get_landmarks(self, frame):
        """Extract landmarks from a frame using Tasks API."""
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)
        
        detection_result = self.detector.detect(mp_image)
        
        if not detection_result.face_landmarks:
            return None
        
        return detection_result.face_landmarks[0]

    def normalize_eye(self, frame, landmarks, eye_side='left', target_size=(64, 32), method='new'):
        """
        Supports two methods:
        - 'old': Raw Affine Warp + Grayscale.
        - 'new': Affine Warp + Median Blur + CLAHE (Huy's 4-Step).
        """
        h, w, _ = frame.shape
        
        if eye_side == 'left':
            indices = self.LEFT_CORNERS
        else:
            indices = self.RIGHT_CORNERS
            
        p1 = np.array([landmarks[indices[0]].x * w, landmarks[indices[0]].y * h])
        p2 = np.array([landmarks[indices[1]].x * w, landmarks[indices[1]].y * h])
        
        center = (p1 + p2) / 2
        dx, dy = p2 - p1
        angle = np.degrees(np.arctan2(dy, dx))
        
        dist = np.linalg.norm(p2 - p1)
        scale = (target_size[0] * 0.7) / (dist + 1e-6)
        
        M = cv2.getRotationMatrix2D(tuple(center), angle, scale)
        M[0, 2] += (target_size[0] / 2) - center[0]
        M[1, 2] += (target_size[1] / 2) - center[1]
        
        # Warp
        normalized = cv2.warpAffine(frame, M, target_size, flags=cv2.INTER_CUBIC)
        
        if len(normalized.shape) == 3:
            normalized = cv2.cvtColor(normalized, cv2.COLOR_BGR2GRAY)

        if method == 'old':
            return normalized, angle

        # Huy's 4-Step Robustness (Method: 'new')
        # Step 2: Median Blur
        normalized = cv2.medianBlur(normalized, 3)

        # Step 3: Apply Tuned CLAHE
        normalized = self.clahe.apply(normalized)
            
        return normalized, angle

    def rotate_gaze(self, gaze_3d, angle_deg):
        """
        Rotate 3D gaze vector around Z-axis (camera axis).
        """
        angle_rad = np.radians(angle_deg)
        c, s = np.cos(angle_rad), np.sin(angle_rad)
        
        # Rotation matrix around Z
        R = np.array([
            [c, -s, 0],
            [s,  c, 0],
            [0,  0, 1]
        ])
        
        return np.dot(R, gaze_3d)

    def gaze_3d_to_mag(self, gaze_3d):
        """Convert 3D gaze to pitch and yaw (radians)."""
        x, y, z = gaze_3d
        pitch = np.arcsin(-y)
        yaw = np.arctan2(-x, -z)
        return np.array([pitch, yaw])

    def extract_patches(self, eye_img, patch_size=8):
        """
        Extract 4 patches using high-quality Interpolation (CUBIC).
        """
        h, w = eye_img.shape
        patches = []
        
        step_w = w // 2
        step_h = h // 2
        
        for i in range(2):
            for j in range(2):
                roi = eye_img[i*step_h : (i+1)*step_h, j*step_w : (j+1)*step_w]
                # Using INTER_CUBIC instead of INTER_AREA for better upsampling (Huy's Suggestion)
                patch = cv2.resize(roi, (patch_size, patch_size), interpolation=cv2.INTER_CUBIC)
                patches.append(patch)
                
        return np.array(patches)

def preprocess_frame(frame, preprocessor, patch_size=8):
    landmarks = preprocessor.get_landmarks(frame)
    if landmarks is None:
        return None
    
    left_eye, _ = preprocessor.normalize_eye(frame, landmarks, 'left')
    right_eye, _ = preprocessor.normalize_eye(frame, landmarks, 'right')
    
    left_patches = preprocessor.extract_patches(left_eye, patch_size)
    right_patches = preprocessor.extract_patches(right_eye, patch_size)
    
    return {
        'left': left_patches,
        'right': right_patches,
        'landmarks': landmarks
    }