#!/usr/bin/env python3
"""
Test pitch and roll angle calculations to investigate issues
"""

import sys
import os

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

try:
    from src.utils.pose_detection import PoseDetector, RETINAFACE_AVAILABLE
    import sqlite3
    from pathlib import Path
    
    def test_retinaface_landmarks():
        """Test what landmarks RetinaFace actually provides"""
        if not RETINAFACE_AVAILABLE:
            print("❌ RetinaFace not available")
            return
        
        print("=" * 60)
        print("TESTING RETINAFACE LANDMARKS")
        print("=" * 60)
        
        # Try to find a test image
        test_image_paths = [
            "demo_photos",
            "data/uploads",
            "data"
        ]
        
        detector = PoseDetector()
        test_image = None
        
        for path in test_image_paths:
            if os.path.exists(path):
                for ext in ['.jpg', '.jpeg', '.png']:
                    for img_file in Path(path).glob(f'*{ext}'):
                        test_image = str(img_file)
                        break
                if test_image:
                    break
        
        if not test_image:
            print("❌ No test image found")
            return
        
        print(f"Using test image: {test_image}")
        
        # Detect faces
        faces = detector.detect_faces_with_landmarks(test_image)
        
        if not faces:
            print("❌ No faces detected")
            return
        
        print(f"\n✅ Found {len(faces)} face(s)")
        
        for face_key, face_data in faces.items():
            print(f"\n{face_key}:")
            landmarks = face_data.get('landmarks', {})
            print(f"  Landmarks keys: {list(landmarks.keys())}")
            
            for landmark_name, position in landmarks.items():
                print(f"    {landmark_name}: {position}")
            
            # Test calculations
            yaw = detector.calculate_yaw_from_landmarks(landmarks)
            pitch = detector.calculate_pitch_from_landmarks(landmarks)
            roll = detector.calculate_roll_from_landmarks(landmarks)
            
            print(f"\n  Calculated angles:")
            print(f"    Yaw:   {yaw:.2f}°" if yaw is not None else "    Yaw:   None")
            print(f"    Pitch: {pitch:.2f}°" if pitch is not None else "    Pitch: None")
            print(f"    Roll:  {roll:.2f}°" if roll is not None else "    Roll:  None")
            
            # Check which landmarks are missing for pitch
            required_for_pitch = ['left_eye', 'right_eye', 'left_mouth', 'right_mouth', 'nose']
            missing = [lm for lm in required_for_pitch if lm not in landmarks]
            if missing:
                print(f"  ⚠️  Missing landmarks for pitch: {missing}")
            
            # Check roll calculation
            if roll is not None:
                left_eye = landmarks.get('left_eye')
                right_eye = landmarks.get('right_eye')
                if left_eye and right_eye:
                    dx = right_eye[0] - left_eye[0]
                    dy = right_eye[1] - left_eye[1]
                    print(f"  Roll calculation details:")
                    print(f"    dx (right_eye[0] - left_eye[0]): {dx:.2f}")
                    print(f"    dy (right_eye[1] - left_eye[1]): {dy:.2f}")
                    print(f"    atan2(dy, dx) = {roll:.2f}°")
                    
                    # Normalize to [-90, 90] range
                    normalized_roll = roll
                    if normalized_roll > 90:
                        normalized_roll = normalized_roll - 180
                    elif normalized_roll < -90:
                        normalized_roll = normalized_roll + 180
                    print(f"    Normalized to [-90, 90]: {normalized_roll:.2f}°")
            
            pose_mode = detector.classify_pose_mode(yaw, pitch, roll)
            print(f"  Pose mode: {pose_mode}")
    
    def analyze_database_angles():
        """Analyze angles in database to find patterns"""
        db_path = "data/punimtag.db"
        
        if not os.path.exists(db_path):
            print(f"❌ Database not found: {db_path}")
            return
        
        print("\n" + "=" * 60)
        print("ANALYZING DATABASE ANGLES")
        print("=" * 60)
        
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        # Get faces with angle data
        cursor.execute("""
            SELECT id, pose_mode, yaw_angle, pitch_angle, roll_angle
            FROM faces
            WHERE yaw_angle IS NOT NULL OR pitch_angle IS NOT NULL OR roll_angle IS NOT NULL
            LIMIT 20
        """)
        
        faces = cursor.fetchall()
        print(f"\nFound {len(faces)} faces with angle data\n")
        
        for face in faces:
            print(f"Face ID {face['id']}: {face['pose_mode']}")
            print(f"  Yaw:   {face['yaw_angle']:.2f}°" if face['yaw_angle'] else "  Yaw:   None")
            print(f"  Pitch: {face['pitch_angle']:.2f}°" if face['pitch_angle'] else "  Pitch: None")
            print(f"  Roll:  {face['roll_angle']:.2f}°" if face['roll_angle'] else "  Roll:  None")
            
            # Check roll normalization
            if face['roll_angle'] is not None:
                roll = face['roll_angle']
                normalized = roll
                if normalized > 90:
                    normalized = normalized - 180
                elif normalized < -90:
                    normalized = normalized + 180
                print(f"  Roll normalized: {normalized:.2f}°")
            print()
        
        conn.close()
    
    if __name__ == "__main__":
        test_retinaface_landmarks()
        analyze_database_angles()
        
except ImportError as e:
    print(f"❌ Import error: {e}")
    print("Make sure you're in the project directory and dependencies are installed")

