import cv2
import mediapipe as mp
import numpy as np

class DanceVisualizer:
    def __init__(self):
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_pose = mp.solutions.pose
        
    def draw_pose(self, frame, pose_landmarks):
        """Draw pose landmarks on frame"""
        if pose_landmarks is None:
            return frame
            
        # Create copy of frame
        annotated_frame = frame.copy()
        
        # Draw the pose landmarks
        self.mp_drawing.draw_landmarks(
            annotated_frame,
            pose_landmarks,
            self.mp_pose.POSE_CONNECTIONS,
            self.mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2),
            self.mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2)
        )
        
        return annotated_frame