import pygame

class SquareSpriteSheet():
    def __init__(self, filename, size, numberOfImages, transparency_color=(0,0,0)) -> None:
        self.sheet = pygame.image.load(filename).convert_alpha()
        self.size = size
        self.numberOfImages = numberOfImages
        self.transparency_color = transparency_color
        self.images = []

        self.load_images()

    def load_images(self):
        for i in range(self.numberOfImages):
            image = pygame.Surface((self.size, self.size)).convert_alpha()
            image.blit(self.sheet, (0,0), ((i * self.size), 0, self.size, self.size))
            image.set_colorkey(self.transparency_color)
            self.images.append(image)

    def get_image_by_frame(self, frame, scale=1, angle=0):
            image = self.images[frame]
            image = pygame.transform.scale(image, (self.size * scale, self.size * scale))
            image = self.rot_center(image, angle)
            image.set_colorkey(self.transparency_color)
            return image

    def rot_center(self, image, angle):
        """rotate an image while keeping its center and size"""
        orig_rect = image.get_rect()
        rot_image = pygame.transform.rotate(image, angle)
        rot_rect = orig_rect.copy()
        rot_rect.center = rot_image.get_rect().center
        rot_image = rot_image.subsurface(rot_rect).copy()
        return rot_image