import numpy as np
from sign_language_tools.visualization.video.displayable import Displayable
from sign_language_tools.visualization.video.cv.landmarks import draw_pose


class Poses(Displayable):

    def __init__(
            self,
            resolution: tuple[int, int],
    ):
        self.resolution = resolution
        self.n_poses = 0
        self.poses = {}

    def add_pose(
            self,
            name: str,
            pose_sequence: np.ndarray,
            edges: tuple[tuple[int, int]],
            show_vertices: bool,
            vertex_color: tuple[int, int, int],
            edge_color: tuple[int, int, int],
    ):
        # (T, V, D) or (T, VxD)
        if len(pose_sequence.shape) == 2:
            pose_sequence = pose_sequence.reshape((pose_sequence.shape[0], -1, 2))

        self.n_poses = pose_sequence.shape[0]
        self.poses[name] = {
            'vertices': pose_sequence,
            'edges': edges,
            'show_vertices': show_vertices,
            'vertex_color': vertex_color,
            'edge_color': edge_color,
        }

    def get_img(self, frame_number: int) -> np.ndarray:
        width, height = self.resolution
        img = np.zeros((height, width, 3), dtype='uint8')
        self.draw(img, frame_number)
        return img

    def draw(self, img: np.ndarray, frame_number: int):
        for name in self.poses:
            pose = self.poses[name]
            draw_pose(
                img,
                pose['vertices'][frame_number],
                pose['edges'],
                show_vertices=pose['show_vertices'],
                vertex_color=pose['vertex_color'],
                edge_color=pose['edge_color'],
            )
