import torch
import cv2
import numpy as np
import sys
import os
from pathlib import Path

# Add project root to path
sys.path.append(str(Path(__file__).parent.parent.parent))
from src.models.teacher import load_teacher_model
from src.utils.preprocess import GazePreprocessor

def get_face_crop(frame, landmarks, target_size=(224, 224)):
    h, w, _ = frame.shape
    coords = np.array([[lm.x * w, lm.y * h] for lm in landmarks])
    min_x, min_y = np.min(coords, axis=0)
    max_x, max_y = np.max(coords, axis=0)
    width = max_x - min_x
    height = max_y - min_y
    center_x = (min_x + max_x) / 2
    center_y = (min_y + max_y) / 2
    size = max(width, height) * 1.5
    x1 = int(max(0, center_x - size / 2))
    y1 = int(max(0, center_y - size / 2))
    x2 = int(min(w, center_x + size / 2))
    y2 = int(min(h, center_y + size / 2))
    face_img = frame[y1:y2, x1:x2]
    if face_img.size == 0: return None
    face_img = cv2.resize(face_img, target_size)
    face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
    face_img = face_img.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    face_img = (face_img - mean) / std
    face_img = np.transpose(face_img, (2, 0, 1))
    return face_img, (x1, y1, x2, y2)

def draw_gaze(image, pitch, yaw, length=100, color=(0, 255, 0)):
    # Simple 2D projection for visualization
    # Pitch (up/down), Yaw (left/right)
    # We assume gaze starts at center
    h, w = image.shape[:2]
    cx, cy = w // 2, h // 2
    
    # Negative yaw because +x is right, but gaze x is usually -x
    dx = -length * np.sin(np.radians(yaw))
    dy = -length * np.sin(np.radians(pitch)) # -pitch because +y is down
    
    cv2.line(image, (cx, cy), (int(cx + dx), int(cy + dy)), color, 3)
    return image

def visualize_teacher():
    checkpoint_path = 'checkpoints/resnet50.pt'
    # Use absolute path for MediaPipe model
    model_path = os.path.abspath('face_landmarker.task')
    
    device = 'cpu'
    model = load_teacher_model(checkpoint_path, backbone='resnet50', device=device)
    preprocessor = GazePreprocessor(model_path=model_path)
    
    # Try multiple images from p00/day01
    img_dir = 'data/MPIIGaze/MPIIGaze/MPIIGaze/Data/Original/p00/day01'
    for i in range(1, 10):
        img_path = os.path.join(img_dir, f"{i:04d}.jpg")
        if not os.path.exists(img_path): continue
        
        print(f"Testing on {img_path}...")
        frame = cv2.imread(img_path)
        landmarks = preprocessor.get_landmarks(frame)
        if landmarks is None:
            print(f"No landmarks found for {img_path}.")
            continue
            
        face_input, (x1, y1, x2, y2) = get_face_crop(frame, landmarks)
        input_tensor = torch.from_numpy(face_input).unsqueeze(0).to(device)
        
        with torch.no_grad():
            p_logits, y_logits = model(input_tensor)
            p_deg, y_deg = model.get_angles(p_logits, y_logits)
            
        print(f"Predicted: Pitch {p_deg.item():.2f}, Yaw {y_deg.item():.2f}")
        
        # Visualize
        face_vis = frame[y1:y2, x1:x2].copy()
        face_vis = draw_gaze(face_vis, p_deg.item(), y_deg.item())
        
        output_path = f'data/verification/teacher_test_{i}.png'
        cv2.imwrite(output_path, face_vis)
        print(f"Saved to {output_path}")
        break

if __name__ == '__main__':
    visualize_teacher()