jogai/jogai/group.py

from random import randint
from typing import Type

import pygame

from jogai.character import Character
from jogai.translations import gettext as _


class Group(pygame.sprite.Group):
    """
    Generic group of sprites.

    Examples:
        >>> g = Group()
    """

    def __init__(self, *args, **kwgargs):
        super().__init__(*args, **kwgargs)


class RandomGroup(Group):
    """
    Group of randomly distributed characters.

    Examples:
        >>> protagonist = Character('protagonist.png')
        >>> rg = RandomGroup(protagonist, 5, (100, 500), (0, 300))
    """

    def __init__(
        self,
        model: Type['Character'],
        amount: int,
        xrange: tuple,
        yrange: tuple,
    ):
        # TODO: check model type
        if not isinstance(amount, int):
            raise TypeError(
                _('The amount of characters must be a integer greater than 1.')
            )
        if amount < 2:
            raise ValueError(
                _('The amount of characters must be a integer greater than 1.')
            )

        if isinstance(xrange, list):
            xrange = tuple(xrange)
        if isinstance(yrange, list):
            yrange = tuple(yrange)
        if not isinstance(xrange, tuple) or not isinstance(yrange, tuple):
            raise TypeError(_('Both x and y ranges must be tuples.'))
        if len(xrange) != 2 or len(yrange) != 2:
            raise ValueError(
                _('Both x and y ranges must have only two items (min, max)')
            )
        if xrange[0] > xrange[1] or yrange[0] > yrange[1]:
            raise ValueError(
                _('Both x and y ranges must be sorted (min, max)')
            )

        self.amount = amount
        self.xrange = xrange
        self.yrange = yrange

        characters_list = []
        for pos in range(self.amount):
            new_character = model.copy()
            new_character.set_position(
                (
                    randint(*xrange),
                    randint(*yrange),
                )
            )
            characters_list.append(new_character)

        super().__init__(characters_list)


class RectangleGroup(Group):
    """
    Group of characters forming a rectangle.

    Examples:
        >>> protagonist = Character('protagonist.png')
        >>> rg = RectangleGroup(protagonist, 5, 3)
    """

    def __init__(self, model: Type['Character'], base: int, height: int):
        # TODO: check model type
        if not isinstance(base, int) or not isinstance(height, int):
            raise TypeError(_('Both the base and rectangle must be integers.'))
        if base < 1 or height < 1:
            raise ValueError(
                _('Both the base and rectangle must be positive.')
            )
        if base == 1 and height == 1:
            raise ValueError(
                _(
                    'At least one of the dimensions of the rectangle must be greater than 1.'
                )
            )
        self.base = base
        self.height = height
        self.amount = base * height

        characters_list = []
        for row in range(self.height):
            for col in range(self.base):
                if not characters_list:
                    new_character = model
                else:
                    new_character = model.copy()
                new_character.set_position(
                    (
                        model.x + model.width * col,
                        model.y + model.height * row,
                    )
                )
                characters_list.append(new_character)
        super().__init__(characters_list)


class SquareGroup(RectangleGroup):
    """
    Group of characters forming a square.

    Examples:
        >>> protagonist = Character('protagonist.png')
        >>> sg = SquareGroup(protagonist, 3)
    """

    def __init__(self, model: Type['Character'], side: int):
        # TODO: check model type
        if not isinstance(side, int):
            raise TypeError(_('The side of the square must be an integer.'))
        if side < 2:
            raise ValueError(
                _('The side of the square must be greater than 1.')
            )
        self.side = side
        super().__init__(model, side, side)


class StairsGroup(Group):
    """
    Group of characters forming stairs.

    Examples:
        >>> protagonist = Character('protagonist.png')
        >>> sg = StairsGroup(protagonist, 3)
    """

    def __init__(
        self,
        model: Type['Character'],
        steps: int,
        horizontal_flip: bool = False,
        vertical_flip: bool = False,
    ):
        # TODO: check model type and adjust the position of elements
        if not isinstance(steps, int):
            raise TypeError(_('The steps of the stairs must be an integer.'))
        if steps < 2:
            raise ValueError(
                _('The steps of the stairs must be greater than 1.')
            )
        if not isinstance(horizontal_flip, bool) or not isinstance(
            vertical_flip, bool
        ):
            raise TypeError(
                _('Both horizontal and vertical flips must be a boolean.')
            )

        self.steps = steps
        self.amount = steps * (steps + 1) // 2
        self.horizontal_flip = horizontal_flip
        self.vertical_flip = vertical_flip

        horizontal_sign = 1 if horizontal_flip else -1
        vertical_sign = -1 if vertical_flip else 1
        characters_list = [model]
        for row in range(2, self.steps + 1):
            for col in range(row):
                new_character = model.copy()
                new_character.name = model.name
                new_character.set_position(
                    (
                        model.x + model.width * col * horizontal_sign,
                        model.y + model.height * (row - 1) * vertical_sign,
                    )
                )
                characters_list.append(new_character)
        super().__init__(characters_list)