"""
PriviGaze Inference Script - Run the student model on-device

Usage:
    python inference.py --model ./checkpoints/student_best.pt --image face.jpg
"""
import argparse
import torch
import numpy as np
from PIL import Image, ImageOps
from models.student import PriviGazeStudent


def load_model(checkpoint_path, device='cpu'):
    model = PriviGazeStudent()
    ckpt = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(ckpt.get('student_state_dict', ckpt))
    model.eval().to(device)
    return model


def preprocess(image_path, size=224):
    img = Image.open(image_path).convert('L')
    img = img.resize((size, size))
    arr = np.array(img).astype(np.float32) / 255.0
    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)
    return tensor * 2 - 1


def predict(model, tensor, device='cpu'):
    tensor = tensor.to(device)
    with torch.no_grad():
        pitch, yaw, _ = model(tensor)
    return pitch.item(), yaw.item()


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--model', type=str, required=True, help='Path to student checkpoint')
    p.add_argument('--image', type=str, required=True, help='Path to grayscale face image')
    p.add_argument('--device', type=str, default='cpu')
    args = p.parse_args()

    device = torch.device(args.device)
    model = load_model(args.model, device)
    tensor = preprocess(args.image)
    pitch, yaw = predict(model, tensor, device)
    print(f"Gaze: pitch={pitch:+.2f}deg, yaw={yaw:+.2f}deg")
    print(f"Direction: {'up' if pitch>0 else 'down'} {abs(pitch):.1f}deg, "
          f"{'right' if yaw>0 else 'left'} {abs(yaw):.1f}deg")


if __name__ == "__main__":
    main()