import numpy as np

PIECES = {
    "O": np.array([[1, 1], [1, 1]]),  # 2x2 block
    "I": np.array([[1, 1, 1, 1]]),  # 4x1 horizontal bar
    "L": np.array([[1, 0], [1, 0], [1, 1]]),  # L-shaped
}


def check_placement(piece, board, i, j):
    n, m = piece.shape
    if i + n > board.shape[0] or j + m > board.shape[1]:
        return False  # Out of bounds

    # Check for overlap
    for x in range(n):
        for y in range(m):
            if piece[x, y] == 1 and board[i + x, j + y] == 1:
                return False
    return True


def place_piece(piece, board, i, j):
    board[i : i + piece.shape[0], j : j + piece.shape[1]] += piece
    return board


def attempt_rotate(piece, board, i, j):
    rotated_piece = np.rot90(piece)
    if check_placement(rotated_piece, board, i, j):
        return rotated_piece
    return piece


def clear_rows(board):
    clear_indices = np.where(np.all(board == 1, axis=1))[0]
    if clear_indices.size > 0:
        board = np.delete(board, clear_indices, axis=0)
        new_rows = np.zeros((len(clear_indices), board.shape[1]))
        board = np.vstack([new_rows, board])
    return board


def calculate_shadow(piece, board, i, j):
    """Calculate the shadow (landing) position for the current piece."""
    while check_placement(piece, board, i + 1, j):
        i += 1
    return i


def display_board(board, piece, i, j):
    display = board.copy()
    shadow_i = calculate_shadow(piece, board, i, j)
    n, m = piece.shape
    # Overlay shadow
    for x in range(n):
        for y in range(m):
            if piece[x, y] == 1:
                display[shadow_i + x, j + y] = 2  # Shadow with different marker

    # Overlay piece
    for x in range(n):
        for y in range(m):
            if piece[x, y] == 1:
                display[i + x, j + y] = 1
    print(
        "\n".join(
            "".join("#" if x == 1 else "." if x == 0 else "*" for x in row)
            for row in display
        )
    )
    print()


def main():
    board = np.zeros((20, 10), dtype=int)
    piece_types = list(PIECES.keys())
    current_piece = PIECES[np.random.choice(piece_types)]
    START_POS = (0, 4)  # Starting position
    i, j = START_POS
    auto_gravity = True  # Auto-gravity enabled by default

    while True:
        display_board(board, current_piece, i, j)
        command = input(
            "Enter command (a=left, d=right, s=down, w=rotate, S=drop, G=toggle gravity, Q=quit): "
        )
        if command == "Q":
            break

        if command == "a" and j > 0 and check_placement(current_piece, board, i, j - 1):
            j -= 1
        elif (
            command == "d"
            and j + current_piece.shape[1] < board.shape[1]
            and check_placement(current_piece, board, i, j + 1)
        ):
            j += 1
        elif command == "w":
            new_piece = attempt_rotate(current_piece, board, i, j)
            if check_placement(new_piece, board, i, j):
                current_piece = new_piece
        elif command == "s":
            if check_placement(current_piece, board, i + 1, j):
                i += 1
            else:
                board = place_piece(current_piece, board, i, j)
                board = clear_rows(board)
                current_piece = PIECES[np.random.choice(piece_types)]
                i, j = 0, 4
        elif command == "S":
            i = calculate_shadow(current_piece, board, i, j)
            board = place_piece(current_piece, board, i, j)
            board = clear_rows(board)
            current_piece = PIECES[np.random.choice(piece_types)]
            i, j = START_POS
        elif command == "G":
            auto_gravity = not auto_gravity

        if auto_gravity:
            if check_placement(current_piece, board, i + 1, j):
                i += 1
            else:
                board = place_piece(current_piece, board, i, j)
                board = clear_rows(board)
                current_piece = PIECES[np.random.choice(piece_types)]
                i, j = START_POS


if __name__ == "__main__":
    main()