import cv2
import time
import argparse
from pose_analyzer import PoseAnalyzer

def process_video(video_source, analyzer):
    # Initialize video capture
    cap = cv2.VideoCapture(video_source)
    
    # Set window properties
    cv2.namedWindow('Bodybuilding Pose Analyzer', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Bodybuilding Pose Analyzer', 1280, 720)
    
    # FPS calculation variables
    prev_time = 0
    curr_time = 0
    
    while cap.isOpened():
        # Read frame
        ret, frame = cap.read()
        if not ret:
            break
            
        # Calculate FPS
        curr_time = time.time()
        fps = 1 / (curr_time - prev_time) if prev_time > 0 else 0
        prev_time = curr_time
        
        # Process frame
        frame_with_pose, analysis = analyzer.process_frame(frame)
        
        # Add FPS and analysis text to frame
        cv2.putText(frame_with_pose, f'FPS: {fps:.1f}', (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # Display feedback
        if 'error' not in analysis:
            y_offset = 70
            cv2.putText(frame_with_pose, f'Pose: {analysis["pose_type"]}', (10, y_offset),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            for angle_name, angle_value in analysis['angles'].items():
                y_offset += 40
                cv2.putText(frame_with_pose, f'{angle_name}: {angle_value:.1f}°', (10, y_offset),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            for correction in analysis['corrections']:
                y_offset += 40
                cv2.putText(frame_with_pose, correction, (10, y_offset),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        else:
            cv2.putText(frame_with_pose, analysis['error'], (10, 70),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        # Display the frame
        cv2.imshow('Bodybuilding Pose Analyzer', frame_with_pose)
        
        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # Release resources
    cap.release()
    cv2.destroyAllWindows()

def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Bodybuilding Pose Analyzer Demo')
    parser.add_argument('--video', type=str, help='Path to video file (optional)')
    args = parser.parse_args()
    
    # Initialize the pose analyzer
    analyzer = PoseAnalyzer()
    
    # Process video (either webcam or file)
    video_source = args.video if args.video else 0
    process_video(video_source, analyzer)

if __name__ == '__main__':
    main()