nmreval/src/gui_qt/lib/stuff.py

import random
import time

import numpy as np

from ..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 = 125
    WIDTH = 35
    HEIGHT = 35

    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.next_direction = []

        self._direction_pairs = {'l': 'r', 'u': 'd', 'r': 'l', 'd': 'u'}

        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.next_direction:
            turn = self.next_direction.pop()
            if self.direction != self._direction_pairs[turn]:
                self.direction = turn

        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) in self.snake:
            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):
            self.next_direction.append('l')

        elif key in (QtCore.Qt.Key_Right, QtCore.Qt.Key_D):
            self.next_direction.append('r')

        elif key in (QtCore.Qt.Key_Down, QtCore.Qt.Key_S):
            self.next_direction.append('d')

        elif key in (QtCore.Qt.Key_Up, QtCore.Qt.Key_W):
            self.next_direction.append('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'


class Field(QtWidgets.QToolButton):
    NUM_COLORS = {
        1: '#1e46a4',
        2: '#f28e2b',
        3: '#e15759',
        4: '#76b7b2',
        5: '#59a14f',
        6: '#edc948',
        7: '#b07aa1',
        8: '#ff9da7',
        'X': '#ff0000',
    }

    flag_change = QtCore.pyqtSignal(bool)

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

        self.setFixedSize(QtCore.QSize(30, 30))
        f = self.font()
        f.setPointSize(24)
        f.setWeight(75)
        self.setFont(f)

        self.x = x
        self.y = y

        self.neighbors = []
        self.grid = self.parent()

        _size = self.grid.gridsize
        for x_i in range(max(0, self.x-1), min(self.x+2, _size[0])):
            for y_i in range(max(0, self.y-1), min(self.y+2, _size[1])):
                if (x_i, y_i) == (self.x, self.y):
                    continue
                self.neighbors.append((x_i, y_i))

        self.is_mine = False
        self.is_flagged = False
        self.is_known = False
        self.has_died = False

    def mousePressEvent(self, evt: QtGui.QMouseEvent) -> None:
        if self.grid.status == 'finished':
            return

        if evt.button() == QtCore.Qt.RightButton:
            self.set_flag()
        elif evt.button() == QtCore.Qt.LeftButton:
            self.select()
        else:
            super().mousePressEvent(evt)

    def __str__(self) -> str:
        if self.is_mine:
            return 'X'
        else:
            return str(self.visible_mines)

    @property
    def visible_mines(self) -> int:
        return sum(self.grid.map[x_i][y_i].is_mine for x_i, y_i in self.neighbors)

    @QtCore.pyqtSlot()
    def select(self):
        if self.is_flagged:
            return

        if self.is_mine:
            self.setText('X')
            num = 'X'
            self.has_died = True
            self.setStyleSheet(f'color: {self.NUM_COLORS[num]}')
        else:
            self.is_known = True
            self.setEnabled(False)
            self.setAutoRaise(True)

            num = self.visible_mines
            if num != 0:
                self.setText(str(num))
                self.setStyleSheet(f'color: {self.NUM_COLORS[num]}')
            else:
                self.grid.reveal_neighbors(self)

        self.clicked.emit()

    @QtCore.pyqtSlot()
    def set_flag(self):
        if self.is_flagged:
            self.setText('')
        else:
            self.setText('!')
        self.is_flagged = not self.is_flagged
        self.has_died = False

        self.clicked.emit()
        self.flag_change.emit(self.is_flagged)


class QMines(QtWidgets.QMainWindow):
    LEVELS = {
        'easy': ((9, 9), 10),
        'middle': ((16, 16), 40),
        'hard': ((16, 30), 99),
        'very hard': ((16, 30), 10),
    }

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

        self.map = []

        self.status = 'waiting'
        self.open_fields = 0
        self.num_flags = 0

        self._start = 0

        self.gridsize, self.mines = (1, 1), 1

        self._init_ui()

    def _init_ui(self):
        layout = QtWidgets.QGridLayout()
        layout.setSpacing(0)
        self.central = QtWidgets.QWidget()
        self.setCentralWidget(self.central)

        layout.addItem(QtWidgets.QSpacerItem(0, 0, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding),
                       0, 0)

        self.grid_layout = QtWidgets.QGridLayout()
        self.grid_layout.setSpacing(0)
        self.map_widget = QtWidgets.QFrame()
        self.map_widget.setFrameStyle(2)
        self.map_widget.setLayout(self.grid_layout)
        layout.addWidget(self.map_widget, 1, 1)

        self.new_game = QtWidgets.QPushButton('New game')
        self.new_game.pressed.connect(self._init_map)
        layout.addWidget(self.new_game, 0, 1)

        layout.addItem(QtWidgets.QSpacerItem(0, 0, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding),
                       2, 2)

        self.central.setLayout(layout)

        self.timer = QtCore.QTimer()
        self.timer.setInterval(1000)
        self.timer.timeout.connect(self.update_time)

        self.timer_message = QtWidgets.QLabel('0 s')
        self.statusBar().addWidget(self.timer_message)

        self.mine_message = QtWidgets.QLabel(f'0 / {self.mines}')
        self.statusBar().addWidget(self.mine_message)

        self.dead_message = QtWidgets.QLabel('')
        self.statusBar().addWidget(self.dead_message)

        self.level_cb = QtWidgets.QComboBox()
        self.level_cb.addItems(list(self.LEVELS.keys()))
        self.level_cb.currentTextChanged.connect(self._init_map)
        self.statusBar().addPermanentWidget(self.level_cb)

        self._init_map('easy')

    def _init_map(self, lvl: str = None):
        if lvl is None:
            lvl = self.level_cb.currentText()

        self._clear_map()

        self.gridsize, self.mines = QMines.LEVELS[lvl]
        w, h = self.gridsize
        self.map = [[None] * h for _ in range(w)]

        for x in range(w):
            for y in range(h):
                pos = Field(x, y, parent=self)
                pos.clicked.connect(self.start_game)
                pos.clicked.connect(self.update_status)
                pos.flag_change.connect(self.update_flag)
                self.grid_layout.addWidget(pos, x+1, y+1)
                self.map[x][y] = pos

        self.set_mines()

    def _clear_map(self):
        self.status = 'waiting'
        self.open_fields = 0
        self.num_flags = 0

        self._start = 0

        self.map = []

        while self.grid_layout.count():
            child = self.grid_layout.takeAt(0)
            w = child.widget()
            if w:
                self.grid_layout.removeWidget(w)
                w.deleteLater()
            else:
                self.grid_layout.removeItem(child)

    def set_mines(self):
        count = 0
        w, h = self.gridsize
        while count < self.mines:
            n = random.randint(0, w * h - 1)
            row = n // h
            col = n % h
            pos = self.map[row][col]  # type: Field
            if pos.is_mine:
                continue
            pos.is_mine = True
            count += 1

        self.status = 'waiting'
        self.open_fields = 0
        self.num_flags = 0
        self._start = 0

        self.timer_message.setText('0 s')
        self.mine_message.setText(f'0 / {self.mines}')
        self.dead_message.setText('')

    def reveal_neighbors(self, pos: Field):
        for x_i, y_i in pos.neighbors:
            field_i = self.map[x_i][y_i]  # type: Field
            if field_i.isEnabled():
                field_i.select()

    def start_game(self):
        if self.status == 'waiting':
            self.status = 'running'
            self._start = time.time()
            self.timer.start(1000)

    def update_time(self):
        if self.status == 'running':
            self.timer_message.setText(f'{time.time()-self._start:3.0f} s')

    def update_status(self):
        if self.status == 'finished':
            return

        pos = self.sender()  # type: Field
        if pos.is_known:
            self.open_fields += 1
            if self.open_fields == self.gridsize[0] * self.gridsize[1] - self.mines:
                self.timer.stop()

                _ = QtWidgets.QMessageBox.information(self, 'Game finished', 'Game finished!!!')
                self.status = 'finished'

        elif pos.has_died:
            dead_cnt = self.dead_message.text()
            if dead_cnt == '':
                self.dead_message.setText('(Deaths: 1)')
            else:
                self.dead_message.setText(f'(Deaths: {int(dead_cnt[9:-1])+1})')

    @QtCore.pyqtSlot(bool)
    def update_flag(self, state: bool):
        num_mines = int(self.mine_message.text().split()[0])
        if state:
            num_mines += 1
        else:
            num_mines -= 1

        self.mine_message.setText(f'{num_mines} / {self.mines}')