import numpy as np
import torch

def pitch_yaw_to_vector(pitch_yaw):
    """
    Convert Pitch/Yaw angles (in Radians) to 3D Gaze Vectors.
    pitch_yaw: (N, 2) or (2,)
    """
    if isinstance(pitch_yaw, torch.Tensor):
        pitch = pitch_yaw[..., 0]
        yaw = pitch_yaw[..., 1]
        
        x = -torch.cos(pitch) * torch.sin(yaw)
        y = -torch.sin(pitch)
        z = -torch.cos(pitch) * torch.cos(yaw)
        
        return torch.stack([x, y, z], dim=-1)
    else:
        pitch = pitch_yaw[0]
        yaw = pitch_yaw[1]
        
        x = -np.cos(pitch) * np.sin(yaw)
        y = -np.sin(pitch)
        z = -np.cos(pitch) * np.cos(yaw)
        
        return np.array([x, y, z])

def angular_error(y_pred, y_true):
    """
    Calculate angular error between two pitch/yaw tensors.
    y_pred, y_true: (N, 2) in Radians.
    Returns: Average angular error in Degrees.
    """
    v_pred = pitch_yaw_to_vector(y_pred)
    v_true = pitch_yaw_to_vector(y_true)
    
    # Normalize (just in case)
    v_pred = v_pred / torch.norm(v_pred, dim=-1, keepdim=True)
    v_true = v_true / torch.norm(v_true, dim=-1, keepdim=True)
    
    # Dot product
    cos_sim = torch.sum(v_pred * v_true, dim=-1)
    # Clamp to avoid numerical issues with acos
    cos_sim = torch.clamp(cos_sim, -1.0 + 1e-7, 1.0 - 1e-7)
    
    # Angle in Radians
    angle_rad = torch.acos(cos_sim)
    
    # Convert to Degrees
    angle_deg = angle_rad * (180.0 / np.pi)
    
    return torch.mean(angle_deg)

if __name__ == "__main__":
    # Test
    y_true = torch.tensor([[0.0, 0.0], [0.1, 0.1]])
    y_pred = torch.tensor([[0.0, 0.1], [0.12, 0.12]]) # Slight error
    
    error = angular_error(y_pred, y_true)
    print(f"Angular Error: {error.item():.4f} degrees")