import cv2
import argparse
from movenet_analyzer import MoveNetAnalyzer

def main():
    parser = argparse.ArgumentParser(description='MoveNet Pose Analysis Demo')
    parser.add_argument('--video', type=str, help='Path to video file (optional)')
    parser.add_argument('--model', type=str, default='lightning', choices=['lightning', 'thunder'],
                      help='MoveNet model variant (lightning or thunder)')
    args = parser.parse_args()

    # Initialize the MoveNet analyzer
    analyzer = MoveNetAnalyzer(model_name=args.model)

    # Initialize video capture
    if args.video:
        cap = cv2.VideoCapture(args.video)
    else:
        cap = cv2.VideoCapture(0)  # Use webcam if no video file provided

    if not cap.isOpened():
        print("Error: Could not open video source")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Process frame
        frame_with_pose, analysis = analyzer.process_frame(frame)

        # Display analysis results
        if 'error' not in analysis:
            # Display pose type
            cv2.putText(frame_with_pose, f"Pose: {analysis['pose_type']}", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # Display angles
            y_offset = 60
            for joint, angle in analysis['angles'].items():
                cv2.putText(frame_with_pose, f"{joint}: {angle:.1f}°", 
                           (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                y_offset += 30

            # Display corrections
            for correction in analysis['corrections']:
                cv2.putText(frame_with_pose, correction, 
                           (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                y_offset += 30
        else:
            cv2.putText(frame_with_pose, analysis['error'], 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        # Display the frame
        cv2.imshow('MoveNet Pose Analysis', frame_with_pose)

        # Break the loop if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()