import random


class Tetris:
    W = 10
    H = 20

    PIECES = {
        'I': [[(0,1),(1,1),(2,1),(3,1)], [(2,0),(2,1),(2,2),(2,3)]],
        'O': [[(1,0),(2,0),(1,1),(2,1)]],
        'T': [[(1,0),(0,1),(1,1),(2,1)], [(1,0),(1,1),(2,1),(1,2)], [(0,1),(1,1),(2,1),(1,2)], [(1,0),(0,1),(1,1),(1,2)]],
        'L': [[(2,0),(0,1),(1,1),(2,1)], [(1,0),(1,1),(1,2),(2,2)], [(0,1),(1,1),(2,1),(0,2)], [(0,0),(1,0),(1,1),(1,2)]],
        'J': [[(0,0),(0,1),(1,1),(2,1)], [(1,0),(2,0),(1,1),(1,2)], [(0,1),(1,1),(2,1),(2,2)], [(1,0),(1,1),(1,2),(0,2)]],
        'S': [[(1,0),(2,0),(0,1),(1,1)], [(1,0),(1,1),(2,1),(2,2)]],
        'Z': [[(0,0),(1,0),(1,1),(2,1)], [(2,0),(1,1),(2,1),(1,2)]],
    }

    def __init__(self, agents=None, state=None):
        self.score = 0
        self.over = False
        self.grid = [[0] * self.W for _ in range(self.H)]
        self.rng = random.Random(0)
        self.next_piece = self._rand_piece()
        self._spawn()
        self.last_render = ''

    def _rand_piece(self):
        return self.rng.choice(list(self.PIECES.keys()))

    def _spawn(self):
        self.piece = self.next_piece
        self.next_piece = self._rand_piece()
        self.rot = 0
        self.px = (self.W // 2) - 2
        self.py = 0
        if not self._fits(self.px, self.py, self.rot):
            self.over = True

    def _cells(self, px, py, rot):
        rstates = self.PIECES[self.piece]
        r = rot % len(rstates)
        return [(px + x, py + y) for (x, y) in rstates[r]]

    def _fits(self, px, py, rot):
        for x, y in self._cells(px, py, rot):
            if x < 0 or x >= self.W or y < 0 or y >= self.H:
                return False
            if self.grid[y][x]:
                return False
        return True

    def _lock(self):
        for x, y in self._cells(self.px, self.py, self.rot):
            if 0 <= y < self.H and 0 <= x < self.W:
                self.grid[y][x] = 1
        # clear lines
        newg = [row for row in self.grid if not all(row)]
        cleared = self.H - len(newg)
        if cleared:
            for _ in range(cleared):
                newg.insert(0, [0] * self.W)
            self.grid = newg
            self.score += 100 * cleared
        self._spawn()

    # API methods used by tests / agent
    def step(self, action=0):
        # actions: 0 noop, 1 left, 2 right, 3 rotate, 4 soft drop
        if self.over:
            return
        if action == 1:
            if self._fits(self.px - 1, self.py, self.rot):
                self.px -= 1
        elif action == 2:
            if self._fits(self.px + 1, self.py, self.rot):
                self.px += 1
        elif action == 3:
            if self._fits(self.px, self.py, self.rot + 1):
                self.rot += 1
        elif action == 4:
            if self._fits(self.px, self.py + 1, self.rot):
                self.py += 1
            else:
                self._lock()
            return
        # gravity
        if self._fits(self.px, self.py + 1, self.rot):
            self.py += 1
        else:
            self._lock()

    def is_over(self):
        return bool(self.over)

    def get_score(self):
        return int(self.score)

    # rendering stubs (retro compatibility)
    def clear(self):
        pass

    def draw_board(self):
        lines = []
        occupied = set(self._cells(self.px, self.py, self.rot)) if not self.over else set()
        for y in range(self.H):
            row = ''
            for x in range(self.W):
                if (x, y) in occupied:
                    row += '[]'
                else:
                    row += '##' if self.grid[y][x] else '..'
            lines.append(row)
        self.last_render = '\n'.join(lines)
        return self.last_render

    def draw_piece(self):
        return

    def refresh(self):
        return

    def reset(self):
        self.__init__()