#!/usr/bin/env python
from typing import Any
import numpy as np
import random

BOARD_SIZE = 14


def make_board():
    a = np.array([[0 for i in range(BOARD_SIZE)] for j in range(BOARD_SIZE)])
    a[4, 4] = -1
    a[9, 9] = -1
    return a


tiles = [
    np.array([[1]]),
    np.array([[1], [1]]),
    np.array([[1], [1], [1]]),
    np.array([[1, 0], [1, 1]]),
    np.array([[1], [1], [1], [1]]),
    np.array([[1, 0], [1, 0], [1, 1]]),
    np.array([[1, 0], [1, 1], [1, 0]]),
    np.array([[1, 1], [1, 1]]),
    np.array([[1, 1, 0], [0, 1, 1]]),
    np.array([[1], [1], [1], [1], [1]]),
    np.array([[1, 0], [1, 0], [1, 0], [1, 1]]),
    np.array([[1, 0], [1, 0], [1, 1], [0, 1]]),
    np.array([[1, 0], [1, 1], [1, 1]]),
    np.array([[1, 1], [1, 0], [1, 1]]),
    np.array([[1, 0], [1, 1], [1, 0], [1, 0]]),
    np.array([[0, 1, 0], [0, 1, 0], [1, 1, 1]]),
    np.array([[1, 0, 0], [1, 0, 0], [1, 1, 1]]),
    np.array([[1, 1, 0], [0, 1, 1], [0, 0, 1]]),
    np.array([[1, 0, 0], [1, 1, 1], [0, 0, 1]]),
    np.array([[1, 0, 0], [1, 1, 1], [0, 1, 0]]),
    np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]),
]


def get_permutations(which_tiles: list[int]) -> list[tuple[int, np.ndarray]]:
    permutations: list[tuple[int, np.ndarray]] = []

    for i, tile in enumerate(tiles):
        if i not in which_tiles:
            continue

        rots = [np.rot90(tile, k) for k in range(4)]
        flips = [np.flip(r, axis=1) for r in rots]  # flip horizontally
        all_orients = rots + flips  # 8 orientations

        seen: set[tuple[Any, bytes]] = set()
        for t in all_orients:
            key = (t.shape, t.tobytes())
            if key not in seen:
                seen.add(key)
                permutations.append((i, t))

    return permutations


def can_place(
    board: np.ndarray, tile: np.ndarray, player: int
) -> list[tuple[int, int]]:
    placements: list[tuple[int, int]] = []
    has_minus_one = False
    for x in range(BOARD_SIZE):
        for y in range(BOARD_SIZE):
            if board[x, y] == -1:
                has_minus_one = True
            with np.nditer(tile, flags=["multi_index"]) as it:
                for v in it:
                    if v == 1:
                        (i, j) = it.multi_index
                        if x + i >= BOARD_SIZE:
                            break
                        if y + j >= BOARD_SIZE:
                            break
                        if board[x + i][y + j] > 0:
                            break
                        if x + i - 1 >= 0 and board[x + i - 1][y + j] == player:
                            break
                        if y + j - 1 >= 0 and board[x + i][y + j - 1] == player:
                            break
                        if x + i + 1 < BOARD_SIZE and board[x + i + 1][y + j] == player:
                            break
                        if y + j + 1 < BOARD_SIZE and board[x + i][y + j + 1] == player:
                            break
                else:
                    placements.append((x, y))
    final: list[tuple[int, int]] = []
    if has_minus_one:
        for x, y in placements:
            with np.nditer(tile, flags=["multi_index"]) as it:
                for v in it:
                    (i, j) = it.multi_index
                    if v == 1 and board[x + i, y + j] == -1:
                        final.append((x, y))
                        break
    else:
        for x, y in placements:
            with np.nditer(tile, flags=["multi_index"]) as it:
                for v in it:
                    (i, j) = it.multi_index
                    if (
                        x + i + 1 < BOARD_SIZE
                        and y + j + 1 < BOARD_SIZE
                        and board[x + i + 1][y + j + 1] == player
                    ):
                        final.append((x, y))
                        break
                    if (
                        x + i + 1 < BOARD_SIZE
                        and y + j - 1 >= 0
                        and board[x + i + 1][y + j - 1] == player
                    ):
                        final.append((x, y))
                        break
                    if (
                        x + i - 1 >= 0
                        and y + j + 1 < BOARD_SIZE
                        and board[x + i - 1][y + j + 1] == player
                    ):
                        final.append((x, y))
                        break
                    if (
                        x + i - 1 >= 0
                        and y + j - 1 >= 0
                        and board[x + i - 1][y + j - 1] == player
                    ):
                        final.append((x, y))
                        break
    return final


def do_placement(
    tidx: int,
    tile: np.ndarray,
    placement: tuple[int, int],
    game_state: tuple[np.ndarray, list[int], list[int]],
    player: int,
):
    assert player > 0
    (x, y) = placement
    with np.nditer(tile, flags=["multi_index"]) as it:
        for v in it:
            (i, j) = it.multi_index
            if v == 1:
                game_state[0][x + i, y + j] = player
    game_state[player].remove(tidx)


def print_game_state(game_state: tuple[np.ndarray, list[int], list[int]]):
    (board, p1tiles, p2tiles) = game_state

    for row in board:
        print(
            "".join(
                [
                    " " if x == 0 else "X" if x == 1 else "O" if x == 2 else "S"
                    for x in row
                ]
            )
        )

    print("")
    print(f"Player 1 tiles left: {p1tiles}")
    print(f"Player 2 tiles left: {p2tiles}")


game_state = (
    make_board(),
    [i for i in range(21)],
    [i for i in range(21)],
)


playing = True
player = 1
while playing:
    moves = []
    assert player > 0
    for tidx, tile in get_permutations(game_state[player]):
        for placement in can_place(game_state[0], tile, player):
            moves.append((tidx, tile, placement))

    print_game_state(game_state)
    print(f"player {player} has {len(moves)} options")

    if len(moves) == 0:
        print(f"No moves left, player {player} lost")
        playing = False
        continue

    (tidx, tile, placement) = random.choice(moves)
    do_placement(tidx, tile, placement, game_state, player)

    if player == 1:
        player = 2
    elif player == 2:
        player = 1