import gradio as gr
import cv2
import torch
import numpy as np
from PIL import Image
import os
import sys
from pathlib import Path

# Add src to path
sys.path.append(str(Path(__file__).parent))

from src.models.student import LIPEV2StudentGaze360Gold
from src.utils.preprocess import GazePreprocessor

# 1. Model Initialization
device = "cuda" if torch.cuda.is_available() else "cpu"
model = LIPEV2StudentGaze360Gold()

# Load the best available checkpoint
checkpoint_path = "checkpoints/swa_gold_p11.pt"
if not os.path.exists(checkpoint_path):
    # Fallback search
    pts = list(Path("checkpoints").rglob("*.pt"))
    if pts:
        checkpoint_path = str(pts[0])

if os.path.exists(checkpoint_path):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint)
    print(f"Loaded {checkpoint_path}")

model.to(device)
model.eval()

# 2. Preprocessor Initialization
# Note: face_landmarker.task must be in the root directory for HF Spaces
preprocessor = GazePreprocessor(model_path='face_landmarker.task')

def draw_gaze(frame, pitch, yaw, landmarks):
    h, w, _ = frame.shape
    # Simple visualization logic
    # Calculate eye centers (using indices from preprocessor)
    left_eye_center = np.mean([[landmarks[i].x * w, landmarks[i].y * h] for i in [362, 263]], axis=0).astype(int)
    right_eye_center = np.mean([[landmarks[i].x * w, landmarks[i].y * h] for i in [33, 133]], axis=0).astype(int)
    
    p_rad = np.radians(pitch)
    y_rad = np.radians(yaw)
    
    dx = -np.sin(y_rad) * np.cos(p_rad)
    dy = -np.sin(p_rad)
    
    length = 80
    for center in [left_eye_center, right_eye_center]:
        end_point = (int(center[0] + dx * length), int(center[1] + dy * length))
        cv2.arrowedLine(frame, tuple(center), end_point, (0, 255, 0), 2, tipLength=0.3)
    return frame

def predict(img):
    if img is None:
        return None, "Please upload an image."
    
    # Convert PIL to BGR
    frame = np.array(img)
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    
    # Get landmarks
    landmarks = preprocessor.get_landmarks(frame)
    if not landmarks:
        return img, "No face detected."
    
    # Prepare landmarks (zero-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 [362, 263]], axis=0)
    right_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in [33, 133]], axis=0)
    face_center = (left_c + right_c) / 2
    landmarks_centered = (landmarks_arr - face_center).flatten()
    landmarks_tensor = torch.from_numpy(landmarks_centered).float().unsqueeze(0).to(device)
    
    # Prepare patches (using left eye for inference)
    left_eye_img, _ = preprocessor.normalize_eye(frame, landmarks, 'left')
    patches = preprocessor.extract_patches(left_eye_img, patch_size=16)
    patches_tensor = torch.from_numpy(patches).float().unsqueeze(0).to(device) / 255.0
    
    # Inference
    with torch.no_grad():
        p_logits, y_logits, _ = model(patches_tensor, landmarks_tensor)
        idx = torch.arange(90).float().to(device)
        pitch = (torch.sum(torch.softmax(p_logits, dim=1) * idx, dim=1) * 2 - 90).item()
        yaw = (torch.sum(torch.softmax(y_logits, dim=1) * idx, dim=1) * 2 - 90).item()
    
    # Draw
    res_frame = draw_gaze(frame.copy(), pitch, yaw, landmarks)
    res_img = cv2.cvtColor(res_frame, cv2.COLOR_BGR2RGB)
    
    return Image.fromarray(res_img), f"Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°"

# 3. Gradio Interface
description = """
### LIPE V2: Ultra-Lightweight Gaze Estimation
LIPE (Landmark-guided Image Patch Embedder) is a high-efficiency gaze estimation framework optimized for commodity CPUs.
- **Model Size:** 0.61M Parameters
- **Complexity:** 21.44 MFLOPs
- **Accuracy:** 4.73° (MPII), 7.33° (Gaze360 Zero-shot)
- **Speed:** 30 FPS on Single-threaded Mobile CPU
"""

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(label="Gaze Visualization"), gr.Textbox(label="Result")],
    title="LIPE V2 Live Demo",
    description=description,
    examples=[["data/verification/sample_face.jpg"]] if os.path.exists("data/verification/sample_face.jpg") else None
)

if __name__ == "__main__":
    iface.launch()