nmreval/nmreval/gui_qt/lib/stuff.py

import random
import sys
import numpy as np

from nmreval.gui_qt.Qt import QtWidgets, QtCore, QtGui


__all__ = ['Game']


class Game(QtWidgets.QDialog):
    def __init__(self, mode, parent=None):
        super().__init__(parent=parent)

        layout = QtWidgets.QVBoxLayout()
        layout.setContentsMargins(3, 3, 3, 3)

        self.label = QtWidgets.QLabel(self)
        layout.addWidget(self.label)

        self.startbutton = QtWidgets.QPushButton('Start', self)
        self.startbutton.clicked.connect(self.start)
        layout.addWidget(self.startbutton)

        if mode == 'tetris':
            self._setup_tetris()
        else:
            self._setup_snake()

        layout.addWidget(self.board)
        self.setStyleSheet("""
        Board {
            border: 5px solid black;
        }
        QPushButton {
            font-weight: bold;
        }
        """)

        self.setLayout(layout)

        self.board.new_status.connect(self.new_message)

    def _setup_tetris(self):
        self.board = TetrisBoard(self)
        self.setGeometry(200, 100, 276, 546+self.startbutton.height()+self.label.height())
        self.setWindowTitle('Totally not Tetris')

    def _setup_snake(self):
        self.board = SnakeBoard(self)
        self.setGeometry(200, 100, 406, 406+self.startbutton.height()+self.label.height())
        self.setWindowTitle('Snakey')

    def start(self):

        self.board.start()

    @QtCore.pyqtSlot(str)
    def new_message(self, msg: str):
        self.label.setText(msg)


class Board(QtWidgets.QFrame):
    new_status = QtCore.pyqtSignal(str)

    SPEED = 1000
    WIDTH = 10
    HEIGHT = 10

    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self.next_move)

        self.score = 0
        self._speed = self.SPEED
        self._ispaused = False
        self._isdead = True

        self.setFrameStyle(QtWidgets.QFrame.Box | QtWidgets.QFrame.Raised)
        self.setLineWidth(3)

        self.setFocusPolicy(QtCore.Qt.StrongFocus)
        self.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.MinimumExpanding)

    def _init_game(self):
        raise NotImplementedError

    @property
    def cellwidth(self):
        return int(self.contentsRect().width() // self.WIDTH)

    # square height
    @property
    def cellheight(self):
        return int(self.contentsRect().height() // self.HEIGHT)

    def start(self):
        if self._isdead:
            self._init_game()
            self.new_status.emit(f'Score: {self.score}')
            self.timer.start(self._speed)
            self._isdead = False
            self.setFocus()

    def stop(self, msg):
        self.new_status.emit(f'Score {self.score} // ' + msg)
        self._isdead = True
        self.timer.stop()

    def pause(self):
        if self._ispaused and not self._isdead:
            self.new_status.emit(f'Score {self.score}')
            self.timer.start(self._speed)
        else:
            self.new_status.emit(f'Score {self.score} // Paused')
            self.timer.stop()

        self._ispaused = not self._ispaused

    def next_move(self):
        raise NotImplementedError

    def draw_square(self, painter, x, y, color):
        color = QtGui.QColor(color)
        painter.fillRect(x+1, y+1, self.cellwidth-2, self.cellheight-2, color)

    def draw_circle(self, painter, x, y, color):
        painter.save()

        color = QtGui.QColor(color)
        painter.setPen(QtGui.QPen(color))
        painter.setBrush(QtGui.QBrush(color))
        painter.drawEllipse(x+1, y+1, self.cellwidth, self.cellheight)

        painter.restore()

    def keyPressEvent(self, evt):
        if evt.key() == QtCore.Qt.Key_P:
            self.pause()
        else:
            super().keyPressEvent(evt)


class SnakeBoard(Board):
    SPEED = 100
    WIDTH = 30
    HEIGHT = 30

    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self.snake = [[int(SnakeBoard.WIDTH//2), int(SnakeBoard.HEIGHT//2)],
                      [int(SnakeBoard.WIDTH//2)+1, int(SnakeBoard.HEIGHT//2)]]
        self.current_x_head, self.current_y_head = self.snake[0]
        self.direction = 'l'

        self.food = None
        self.grow_snake = False

    def _init_game(self):
        self.snake = [[int(SnakeBoard.WIDTH//2), int(SnakeBoard.HEIGHT//2)],
                      [int(SnakeBoard.WIDTH//2)+1, int(SnakeBoard.HEIGHT//2)]]
        self.current_x_head, self.current_y_head = self.snake[0]
        self.direction = 'l'

        self.food = None
        self.grow_snake = False
        self.new_food()

    def next_move(self):
        self.snake_move()
        self.got_food()
        self.check_death()
        self.update()

    def snake_move(self):
        if self.direction == 'l':
            self.current_x_head -= 1
        elif self.direction == 'r':
            self.current_x_head += 1

        # y increases top to bottom
        elif self.direction == 'u':
            self.current_y_head -= 1
        elif self.direction == 'd':
            self.current_y_head += 1

        head = [self.current_x_head, self.current_y_head]
        self.snake.insert(0, head)

        if not self.grow_snake:
            self.snake.pop()
        else:
            self.new_status.emit(f'Score: {self.score}')
            self.grow_snake = False

    def got_food(self):
        head = self.snake[0]
        if self.food == head:
            self.new_food()
            self.grow_snake = True
            self.score += 1

    def new_food(self):
        x = random.randint(3, SnakeBoard.WIDTH-3)
        y = random.randint(3, SnakeBoard.HEIGHT-3)

        while [x, y] == self.snake[0]:
            x = random.randint(3, SnakeBoard.WIDTH-3)
            y = random.randint(3, SnakeBoard.HEIGHT-3)

        self.food = [x, y]

    def check_death(self):
        rip_message = ''
        is_dead = False
        if (self.current_x_head < 1) or (self.current_x_head > SnakeBoard.WIDTH-2) or \
                (self.current_y_head < 1) or (self.current_y_head > SnakeBoard.HEIGHT-2):
            rip_message = 'Snake found wall :('
            is_dead = True

        head = self.snake[0]
        for snake_i in self.snake[1:]:
            if snake_i == head:
                rip_message = 'Snake bit itself :('
                is_dead = True
                break

        if is_dead:
            self.stop(rip_message)

    def keyPressEvent(self, event):
        key = event.key()
        if key in (QtCore.Qt.Key_Left, QtCore.Qt.Key_A):
            if self.direction != 'r':
                self.direction = 'l'

        elif key in (QtCore.Qt.Key_Right, QtCore.Qt.Key_D):
            if self.direction != 'l':
                self.direction = 'r'

        elif key in (QtCore.Qt.Key_Down, QtCore.Qt.Key_S):
            if self.direction != 'u':
                self.direction = 'd'

        elif key in (QtCore.Qt.Key_Up, QtCore.Qt.Key_W):
            if self.direction != 'd':
                self.direction = 'u'

        else:
            return super().keyPressEvent(event)

    def paintEvent(self, event):
        painter = QtGui.QPainter(self)

        rect = self.contentsRect()
        boardtop = rect.bottom() - SnakeBoard.HEIGHT * self.cellheight
        boardleft = rect.left()

        for pos in self.snake:
            self.draw_circle(painter,
                             int(boardleft+pos[0]*self.cellwidth),
                             int(boardtop+pos[1]*self.cellheight),
                             'blue')
        if self.food is not None:
            self.draw_square(painter,
                             int(boardleft+self.food[0]*self.cellwidth),
                             int(boardtop+self.food[1]*self.cellheight),
                             'orange')


class TetrisBoard(Board):
    WIDTH = 10
    HEIGHT = 20
    SPEED = 300

    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self._shapes = {
            1: ZShape,
            2: SShape,
            3: Line,
            4: TShape,
            5: Square,
            6: MirrorL
        }

        self._init_game()

    def _init_game(self):
        self.curr_x = 0
        self.curr_y = 0
        self.curr_piece = None
        self.board = np.zeros((TetrisBoard.WIDTH, TetrisBoard.HEIGHT + 2), dtype=int)

    def next_move(self):
        if self.curr_piece is None:
            self.make_new_piece()
        else:
            self.move_down()

    def try_move(self, piece, x, y):
        if piece is None:
            return False

        if x+piece.x.min() < 0 or x+piece.x.max() >= TetrisBoard.WIDTH:
            return False

        if y-piece.y.max() < 0 or y-piece.y.min() >= TetrisBoard.HEIGHT+2:
            return False

        if np.any(self.board[piece.x+x, y-piece.y]) != 0:
            return False

        if piece != self.curr_piece:
            self.curr_piece = piece

        self.curr_x = x
        self.curr_y = y

        self.update()

        return True

    def make_new_piece(self):
        new_piece = self._shapes[random.randint(1, len(self._shapes))]()

        startx = TetrisBoard.WIDTH//2
        starty = TetrisBoard.HEIGHT+2 - 1 + new_piece.y.min()
        if not self.try_move(new_piece, startx, starty):
            self.stop('Game over :(')

    def move_down(self):
        if not self.try_move(self.curr_piece, self.curr_x, self.curr_y-1):
            self.final_destination_reached()

    def drop_to_bottom(self):
        new_y = self.curr_y

        while new_y > 0:
            if not self.try_move(self.curr_piece, self.curr_x, new_y-1):
                break
            new_y -= 1

        self.final_destination_reached()

    def final_destination_reached(self):
        x = self.curr_x+self.curr_piece.x
        y = self.curr_y-self.curr_piece.y
        self.board[x, y] = next(k for k, v in self._shapes.items() if isinstance(self.curr_piece, v))

        self.remove_lines()

        self.curr_piece = None
        self.make_new_piece()

    def remove_lines(self):
        full_rows = np.where(np.all(self.board, axis=0))[0]
        num_rows = len(full_rows)

        if num_rows:
            temp = np.zeros_like(self.board)
            temp[:, :temp.shape[1]-num_rows] = np.delete(self.board, full_rows, axis=1)
            self.board = temp

            self.score += num_rows
            self.new_status.emit(f'Lines: {self.score}')

            if self.score % 10 == 0:
                self._speed += 0.9
                self.timer.setInterval(int(self._speed))

            self.update()

    def keyPressEvent(self, event):
        key = event.key()

        if self.curr_piece is None:
            return super().keyPressEvent(event)

        if key == QtCore.Qt.Key_Left:
            self.try_move(self.curr_piece, self.curr_x-1, self.curr_y)

        elif key == QtCore.Qt.Key_Right:
            self.try_move(self.curr_piece, self.curr_x+1, self.curr_y)

        elif key == QtCore.Qt.Key_Down:
            if not self.try_move(self.curr_piece.rotate(), self.curr_x, self.curr_y):
                self.curr_piece.rotate(clockwise=False)

        elif key == QtCore.Qt.Key_Up:
            if not self.try_move(self.curr_piece.rotate(clockwise=False), self.curr_x, self.curr_y):
                self.curr_piece.rotate()

        elif key == QtCore.Qt.Key_Space:
            self.drop_to_bottom()

        else:
            super().keyPressEvent(event)

    def paintEvent(self, event):
        painter = QtGui.QPainter(self)
        rect = self.contentsRect()
        board_top = rect.bottom() - TetrisBoard.HEIGHT*self.cellheight

        for i in range(TetrisBoard.WIDTH):
            for j in range(TetrisBoard.HEIGHT):
                shape = self.board[i, j]

                if shape:
                    color = self._shapes[shape].color
                    self.draw_square(painter,
                                     rect.left() + i*self.cellwidth,
                                     board_top + (TetrisBoard.HEIGHT-j-1)*self.cellheight, color)

        if self.curr_piece is not None:
            x = self.curr_x + self.curr_piece.x
            y = self.curr_y - self.curr_piece.y

            for i in range(4):
                if TetrisBoard.HEIGHT < y[i]+1:
                    continue

                self.draw_square(painter, rect.left() + x[i] * self.cellwidth,
                                 board_top + (TetrisBoard.HEIGHT-y[i]-1) * self.cellheight,
                                 self.curr_piece.color)


class Tetromino:
    SHAPE = np.array([[0], [0]])
    color = None

    def __init__(self):
        self.shape = self.SHAPE

    def rotate(self, clockwise: bool = True):
        if clockwise:
            self.shape = np.vstack((-self.shape[1], self.shape[0]))
        else:
            self.shape = np.vstack((self.shape[1], -self.shape[0]))

        return self

    @property
    def x(self):
        return self.shape[0]

    @property
    def y(self):
        return self.shape[1]


class ZShape(Tetromino):
    SHAPE = np.array([[0, 0, -1, -1],
                      [-1, 0, 0, 1]])
    color = 'green'


class SShape(Tetromino):
    SHAPE = np.array([[0, 0, 1, 1],
                      [-1, 0, 0, 1]])
    color = 'purple'


class Line(Tetromino):
    SHAPE = np.array([[0, 0, 0, 0],
                      [-1, 0, 1, 2]])
    color = 'red'


class TShape(Tetromino):
    SHAPE = np.array([[-1, 0, 1, 0],
                      [0, 0, 0, 1]])
    color = 'orange'


class Square(Tetromino):
    SHAPE = np.array([[0, 1, 0, 1],
                      [0, 0, 1, 1]])
    color = 'yellow'


class LShape(Tetromino):
    SHAPE = np.array([[-1, 0, 0, 0],
                      [1, -1, 0, 1]])
    color = 'blue'


class MirrorL(Tetromino):
    SHAPE = np.array([[1, 0, 0, 0],
                      [-1, -1, 0, 1]])
    color = 'lightGray'


if __name__ == '__main__':
    app = QtWidgets.QApplication([])
    tetris = Game('snake')
    tetris.show()
    sys.exit(app.exec_())