import cv2
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
from typing import List, Dict, Tuple

class MoveNetAnalyzer:
    KEYPOINT_DICT = {
        'nose': 0,
        'left_eye': 1,
        'right_eye': 2,
        'left_ear': 3,
        'right_ear': 4,
        'left_shoulder': 5,
        'right_shoulder': 6,
        'left_elbow': 7,
        'right_elbow': 8,
        'left_wrist': 9,
        'right_wrist': 10,
        'left_hip': 11,
        'right_hip': 12,
        'left_knee': 13,
        'right_knee': 14,
        'left_ankle': 15,
        'right_ankle': 16
    }

    def __init__(self, model_name="lightning"):
        # Initialize MoveNet model
        if model_name == "lightning":
            self.model = hub.load("https://tfhub.dev/google/movenet/singlepose/lightning/4")
            self.input_size = 192
        else:  # thunder
            self.model = hub.load("https://tfhub.dev/google/movenet/singlepose/thunder/4")
            self.input_size = 256
        
        self.movenet = self.model.signatures['serving_default']
        
        # Define key angles for bodybuilding poses
        self.key_angles = {
            'front_double_biceps': {
                'shoulder_angle': (90, 120),  # Expected angle range
                'elbow_angle': (80, 100),
                'wrist_angle': (0, 20)
            },
            'side_chest': {
                'shoulder_angle': (45, 75),
                'elbow_angle': (90, 110),
                'wrist_angle': (0, 20)
            },
            'back_double_biceps': {
                'shoulder_angle': (90, 120),
                'elbow_angle': (80, 100),
                'wrist_angle': (0, 20)
            }
        }

    def detect_pose(self, frame: np.ndarray, last_valid_landmarks=None) -> Tuple[np.ndarray, List[Dict]]:
        """
        Detect pose in the given frame and return the frame with pose landmarks drawn
        and the list of detected landmarks.
        If detection fails, reuse last valid landmarks if provided.
        """
        # Resize and pad the image to keep aspect ratio
        img = frame.copy()
        img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), self.input_size, self.input_size)
        img = tf.cast(img, dtype=tf.int32)
        
        # Detection
        results = self.movenet(img)
        keypoints = results['output_0'].numpy() # Shape [1, 1, 17, 3]
        
        # Draw the pose landmarks on the frame
        if keypoints[0, 0, 0, 2] > 0.1:  # Lowered confidence threshold for the nose
            # Convert keypoints to image coordinates
            y, x, c = frame.shape
            shaped = np.squeeze(keypoints) # Shape [17, 3]
            
            # Draw keypoints
            for kp in shaped:
                ky, kx, kp_conf = kp
                if kp_conf > 0.1:
                    # Convert to image coordinates
                    x_coord = int(kx * x)
                    y_coord = int(ky * y)
                    cv2.circle(frame, (x_coord, y_coord), 6, (0, 255, 0), -1)
            
            # Convert landmarks to a list of dictionaries
            landmarks = []
            for kp in shaped:
                landmarks.append({
                    'x': float(kp[1]),
                    'y': float(kp[0]),
                    'visibility': float(kp[2])
                })
            
            return frame, landmarks
        
        # If detection fails, reuse last valid landmarks if provided
        if last_valid_landmarks is not None:
            return frame, last_valid_landmarks
        return frame, []

    def calculate_angle(self, landmarks: List[Dict], joint1: int, joint2: int, joint3: int) -> float:
        """
        Calculate the angle between three joints.
        """
        if len(landmarks) < max(joint1, joint2, joint3):
            return None
            
        # Get the coordinates of the three joints
        p1 = np.array([landmarks[joint1]['x'], landmarks[joint1]['y']])
        p2 = np.array([landmarks[joint2]['x'], landmarks[joint2]['y']])
        p3 = np.array([landmarks[joint3]['x'], landmarks[joint3]['y']])
        
        # Calculate the angle
        v1 = p1 - p2
        v2 = p3 - p2
        
        angle = np.degrees(np.arccos(
            np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        ))
        
        return angle

    def analyze_pose(self, landmarks: List[Dict], pose_type: str) -> Dict:
        """
        Analyze the pose and provide feedback based on the pose_type.
        Handles pose_types not in self.key_angles by providing a note.
        """
        feedback = {
            'pose_type': pose_type,
            'angles': {},
            'corrections': [],
            'notes': []  # Initialize notes field
        }

        if not landmarks:
            # If no landmarks, it's a more fundamental issue than just pose_type.
            # The process_frame method already handles this by passing {'error': 'No pose detected'}
            # from self.analyze_pose if landmarks is empty.
            # However, to be safe, if this method is called directly with no landmarks:
            feedback['error'] = 'No landmarks provided for analysis'
            return feedback

        if pose_type not in self.key_angles:
            feedback['notes'].append(f"No specific angle checks defined for pose: {pose_type}")
            # Still return the feedback structure, but angles and corrections will be empty.
            # The 'error' field will not be set here, allowing app.py to distinguish this case.
            return feedback
        
        pose_rules = self.key_angles[pose_type]
        
        if pose_type == 'front_double_biceps':
            # Example: Left Shoulder - Elbow - Wrist for elbow angle
            # Example: Left Hip - Shoulder - Elbow for shoulder angle (arm abduction)
            # Note: These are examples, actual biomechanical definitions can be complex.
            # We'll stick to the previous definition for front_double_biceps shoulder angle for now.
            # Shoulder angle: right_hip - right_shoulder - right_elbow (can also use left)
            # Elbow angle: right_shoulder - right_elbow - right_wrist (can also use left)
            # Wrist angle (simplistic): right_elbow - right_wrist - a point slightly above wrist (not easily done without more points)

            # Using right side for front_double_biceps as an example, consistent with a typical bodybuilding pose display
            # Shoulder Angle (approximating arm abduction/flexion relative to torso)
            # Using Right Hip, Right Shoulder, Right Elbow
            rs = self.KEYPOINT_DICT['right_shoulder']
            re = self.KEYPOINT_DICT['right_elbow']
            rh = self.KEYPOINT_DICT['right_hip']
            rw = self.KEYPOINT_DICT['right_wrist']

            shoulder_angle = self.calculate_angle(landmarks, rh, rs, re)
            if shoulder_angle is not None:
                feedback['angles']['R Shoulder'] = shoulder_angle
                if not (pose_rules['shoulder_angle'][0] <= shoulder_angle <= pose_rules['shoulder_angle'][1]):
                    # Debug print before forming correction string
                    print(f"[MOVENET_DEBUG_CORRECTION] pose_type: {pose_type}, rule_key: 'shoulder_angle', rules_for_angle: {pose_rules.get('shoulder_angle')}")
                    feedback['corrections'].append(
                        f"Adjust R Shoulder to {pose_rules['shoulder_angle'][0]}-{pose_rules['shoulder_angle'][1]} deg"
                    )
            
            elbow_angle = self.calculate_angle(landmarks, rs, re, rw)
            if elbow_angle is not None:
                feedback['angles']['R Elbow'] = elbow_angle
                if not (pose_rules['elbow_angle'][0] <= elbow_angle <= pose_rules['elbow_angle'][1]):
                    feedback['corrections'].append(
                        f"Adjust R Elbow to {pose_rules['elbow_angle'][0]}-{pose_rules['elbow_angle'][1]} deg"
                    )
            # Wrist angle is hard to define meaningfully with current keypoints for this pose, skipping for now.

        elif pose_type == 'side_chest':
            # Assuming side chest often displays left side to judges
            ls = self.KEYPOINT_DICT['left_shoulder']
            le = self.KEYPOINT_DICT['left_elbow']
            lw = self.KEYPOINT_DICT['left_wrist']
            lh = self.KEYPOINT_DICT['left_hip'] # For shoulder angle relative to torso

            # Shoulder angle (e.g. arm flexion/extension in sagittal plane for the front arm)
            # For side chest, the front arm's shoulder angle relative to the torso (hip-shoulder-elbow)
            shoulder_angle = self.calculate_angle(landmarks, lh, ls, le) 
            if shoulder_angle is not None:
                feedback['angles']['L Shoulder'] = shoulder_angle
                if not (pose_rules['shoulder_angle'][0] <= shoulder_angle <= pose_rules['shoulder_angle'][1]):
                    feedback['corrections'].append(
                        f"Adjust L Shoulder to {pose_rules['shoulder_angle'][0]}-{pose_rules['shoulder_angle'][1]} deg"
                    )

            elbow_angle = self.calculate_angle(landmarks, ls, le, lw)
            if elbow_angle is not None:
                feedback['angles']['L Elbow'] = elbow_angle
                if not (pose_rules['elbow_angle'][0] <= elbow_angle <= pose_rules['elbow_angle'][1]):
                    feedback['corrections'].append(
                        f"Adjust L Elbow to {pose_rules['elbow_angle'][0]}-{pose_rules['elbow_angle'][1]} deg"
                    )
            # Wrist angle for side chest is also nuanced, skipping detailed check for now.

        elif pose_type == 'back_double_biceps':
            # Similar to front, but from back. We can calculate for both arms or pick one.
            # Let's do right side for consistency with front_double_biceps example.
            rs = self.KEYPOINT_DICT['right_shoulder']
            re = self.KEYPOINT_DICT['right_elbow']
            rh = self.KEYPOINT_DICT['right_hip']
            rw = self.KEYPOINT_DICT['right_wrist']

            shoulder_angle = self.calculate_angle(landmarks, rh, rs, re)
            if shoulder_angle is not None:
                feedback['angles']['R Shoulder'] = shoulder_angle
                if not (pose_rules['shoulder_angle'][0] <= shoulder_angle <= pose_rules['shoulder_angle'][1]):
                    feedback['corrections'].append(
                        f"Adjust R Shoulder to {pose_rules['shoulder_angle'][0]}-{pose_rules['shoulder_angle'][1]} deg"
                    )
            
            elbow_angle = self.calculate_angle(landmarks, rs, re, rw)
            if elbow_angle is not None:
                feedback['angles']['R Elbow'] = elbow_angle
                if not (pose_rules['elbow_angle'][0] <= elbow_angle <= pose_rules['elbow_angle'][1]):
                    feedback['corrections'].append(
                        f"Adjust R Elbow to {pose_rules['elbow_angle'][0]}-{pose_rules['elbow_angle'][1]} deg"
                    )
        
        # Clear notes if pose_type was valid and processed, unless specific notes were added by pose logic
        if not feedback['notes']: # Only clear if no specific notes were added during pose rule processing
            feedback.pop('notes', None)
        
        return feedback

    def process_frame(self, frame: np.ndarray, pose_type: str = 'front_double_biceps', last_valid_landmarks=None) -> Tuple[np.ndarray, Dict, List[Dict]]:
        """
        Process a single frame, detect pose, and analyze it. Returns frame, analysis, and used landmarks.
        """
        # Detect pose
        frame_with_pose, landmarks = self.detect_pose(frame, last_valid_landmarks=last_valid_landmarks)
        
        # Analyze pose if landmarks are detected
        analysis = self.analyze_pose(landmarks, pose_type) if landmarks else {'error': 'No pose detected'}
        
        return frame_with_pose, analysis, landmarks

    def classify_pose(self, landmarks: List[Dict]) -> str:
        """
        Classify the pose based on keypoint positions and angles.
        Returns one of: 'front_double_biceps', 'side_chest', 'back_double_biceps'.
        """
        if not landmarks or len(landmarks) < 17:
            return 'front_double_biceps'  # Default/fallback

        # Calculate angles for both arms
        # Right side
        rs = self.KEYPOINT_DICT['right_shoulder']
        re = self.KEYPOINT_DICT['right_elbow']
        rh = self.KEYPOINT_DICT['right_hip']
        rw = self.KEYPOINT_DICT['right_wrist']
        # Left side
        ls = self.KEYPOINT_DICT['left_shoulder']
        le = self.KEYPOINT_DICT['left_elbow']
        lh = self.KEYPOINT_DICT['left_hip']
        lw = self.KEYPOINT_DICT['left_wrist']

        # Shoulder angles
        r_shoulder_angle = self.calculate_angle(landmarks, rh, rs, re)
        l_shoulder_angle = self.calculate_angle(landmarks, lh, ls, le)
        # Elbow angles
        r_elbow_angle = self.calculate_angle(landmarks, rs, re, rw)
        l_elbow_angle = self.calculate_angle(landmarks, ls, le, lw)

        # Heuristic rules:
        # - Front double biceps: both arms raised, elbows bent, both shoulders abducted
        # - Side chest: one arm across chest (elbow in front of body), other arm flexed
        # - Back double biceps: both arms raised, elbows bent, but person is facing away (shoulders/hips x order reversed)

        # Use x-coordinates to estimate facing direction
        # If right shoulder x < left shoulder x, assume facing front; else, facing back
        facing_front = landmarks[rs]['x'] < landmarks[ls]['x']

        # Count how many arms are "up" (shoulder angle in expected range)
        arms_up = 0
        if r_shoulder_angle and 80 < r_shoulder_angle < 150:
            arms_up += 1
        if l_shoulder_angle and 80 < l_shoulder_angle < 150:
            arms_up += 1
        elbows_bent = 0
        if r_elbow_angle and 60 < r_elbow_angle < 130:
            elbows_bent += 1
        if l_elbow_angle and 60 < l_elbow_angle < 130:
            elbows_bent += 1

        # Side chest: one arm's elbow is much closer to the body's midline (x of elbow near x of nose)
        nose_x = landmarks[self.KEYPOINT_DICT['nose']]['x']
        le_x = landmarks[le]['x']
        re_x = landmarks[re]['x']
        side_chest_like = (abs(le_x - nose_x) < 0.08 or abs(re_x - nose_x) < 0.08)

        if arms_up == 2 and elbows_bent == 2:
            if facing_front:
                return 'front_double_biceps'
            else:
                return 'back_double_biceps'
        elif side_chest_like:
            return 'side_chest'
        else:
            # Default/fallback
            return 'front_double_biceps'