import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple

@dataclass
class KeyFrame:
    timestamp: float
    landmarks: np.ndarray
    connections: List[Tuple[int, int]]

class AnimationRenderer:
    def __init__(self, fps: int = 30):
        self.fps = fps
        self.keyframes: List[KeyFrame] = []
        self.current_frame = 0
        self.total_frames = 0
        
    def add_keyframe(self, landmarks: np.ndarray, connections: List[Tuple[int, int]], timestamp: float):
        """Add a new keyframe to the animation sequence"""
        keyframe = KeyFrame(timestamp=timestamp, landmarks=landmarks, connections=connections)
        self.keyframes.append(keyframe)
        self.total_frames = max(self.total_frames, int(timestamp * self.fps))
    
    def interpolate_poses(self, start_frame: KeyFrame, end_frame: KeyFrame, alpha: float) -> np.ndarray:
        """Interpolate between two poses using linear interpolation"""
        return start_frame.landmarks + alpha * (end_frame.landmarks - start_frame.landmarks)
    
    def get_frame_at_time(self, time: float) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
        """Get interpolated frame at specified time"""
        if not self.keyframes:
            return None, []
            
        # Find surrounding keyframes
        next_idx = 0
        for i, kf in enumerate(self.keyframes):
            if kf.timestamp > time:
                next_idx = i
                break
                
        if next_idx == 0:
            return self.keyframes[0].landmarks, self.keyframes[0].connections
            
        prev_idx = next_idx - 1
        prev_frame = self.keyframes[prev_idx]
        next_frame = self.keyframes[next_idx]
        
        # Calculate interpolation factor
        alpha = (time - prev_frame.timestamp) / (next_frame.timestamp - prev_frame.timestamp)
        interpolated_landmarks = self.interpolate_poses(prev_frame, next_frame, alpha)
        
        return interpolated_landmarks, prev_frame.connections
    
    def get_next_frame(self) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
        """Get the next frame in the animation sequence"""
        if not self.keyframes or self.current_frame >= self.total_frames:
            return None, []
            
        time = self.current_frame / self.fps
        self.current_frame += 1
        
        return self.get_frame_at_time(time)