#!/usr/bin/env python3
import itertools
from tqdm import tqdm
from dataclasses import dataclass

FILES = "abcdefgh"


@dataclass(frozen=True)
class Square:
    file_num: int
    rank_num: int

    @property
    def rank(self):
        return self.rank_num + 1

    @property
    def file(self):
        return FILES[self.file_num]

    def __add__(self, value: tuple):
        if not (0 <= self.file_num + value[0] < 8
                and 0 <= self.rank_num + value[1] < 8):
            raise ValueError("Square not in board")
        return Square(self.file_num + value[0], self.rank_num + value[1])

    def __sub__(self, value):
        if isinstance(value, Square):
            return (self.file_num - value.file_num,
                    self.rank_num - value.rank_num)
        return self.__add__((-value[0], -value[1]))

    def __str__(self):
        return f"{self.file}{self.rank}"


SQUARES = [Square(f, r) for f in range(8) for r in range(8)]


@dataclass
class Piece:
    letter: str
    colour: str
    square: Square
    moves: tuple

    @property
    def fen(self):
        return self.letter.upper(
        ) if self.colour == "w" else self.letter.lower()

    def get_moves(self, board, captures_only=False):
        for move, rider in self.moves.items():
            for direction in set(
                (a * move[c], b * move[1 - c])
                    for a, b, c in itertools.product((-1, 1), (-1, 1), (0,
                                                                        1))):
                location = self.square
                while True:
                    try:
                        location += direction
                    except ValueError:
                        break
                    square = board[location]
                    if square is None:
                        if not captures_only:
                            yield location
                        if rider:
                            continue
                        break
                    if square.colour != self.colour:
                        yield location
                    break

    def get_moved(self, move):
        return self.__class__(self.colour, move)


class Knight(Piece):
    letter = "N"
    moves = {(1, 2): False}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Bishop(Piece):
    letter = "B"
    moves = {(1, 1): True}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Rook(Piece):
    letter = "R"
    moves = {(0, 1): True}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Queen(Piece):
    letter = "Q"
    moves = {(0, 1): True, (1, 1): True}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Chancellor(Piece):
    letter = "C"
    moves = {(1, 1): True, (1, 2): False}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Marshall(Piece):
    letter = "M"
    moves = {(0, 1): True, (1, 2): False}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


class Pawn(Piece):
    letter = "P"
    moves = None

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)

    def get_moves(self, board, captures_only=False):
        for x in (-1, 1):
            try:
                location = self.square + (x, (1 if self.colour == "w" else -1))
            except ValueError:
                continue
            if (board.en_passant != "-" and location == Square(
                    FILES.index(board.en_passant[0]),
                    int(board.en_passant[1]) - 1,
            )) or (board[location] is not None
                   and board[location].colour != self.colour):
                yield location
        if captures_only:
            return
        location = self.square
        for _ in range(2):
            location += (0, (1 if self.colour == "w" else -1))
            if board[location] is not None:
                break
            yield location
            if not ((self.square.rank == 2 and self.colour == "w") or
                    (self.square.rank == 7 and self.colour == "b")):
                break


class King(Piece):
    letter = "K"
    moves = {(0, 1): False, (1, 1): False}

    def __init__(self, colour, square):
        return super().__init__(self.letter, colour, square, self.moves)


def piece_from_letter(letter, square):
    if letter == " ":
        return None
    return {
        "P": Pawn,
        "N": Knight,
        "B": Bishop,
        "R": Rook,
        "Q": Queen,
        "C": Chancellor,
        "M": Marshall,
        "K": King,
    }[letter.upper()]("w" if letter.isupper() else "b", square)


class Board:
    def __init__(self):
        self.squares = [None for _ in range(64)]
        self.move = "w"
        self.castling = ["K", "Q", "k", "q"]
        self.half_moves = 0
        self.full_moves = 1
        self.en_passant = "-"

    def __getitem__(self, square):
        return self.squares[square.file_num * 8 + square.rank_num]

    def __setitem__(self, square, item):
        self.squares[square.file_num * 8 + square.rank_num] = item

    def to_fen(self):
        fen = []
        for rank in reversed(range(8)):
            rank_text = ""
            blanks = 0
            for file_ in range(8):
                if self[Square(file_, rank)] is None:
                    blanks += 1
                else:
                    if blanks:
                        rank_text += str(blanks)
                    blanks = 0
                    rank_text += self[Square(file_, rank)].fen
            if blanks:
                rank_text += str(blanks)
            fen.append(rank_text)
        fen = "/".join(fen)
        return (
            fen +
            f" {self.move} {''.join(self.castling)} {self.en_passant} {self.half_moves} {self.full_moves}"
            .replace("  ", " - "))

    def is_safe(self):
        return all(
            sum(isinstance(p, King)
                for p in self.get_moved(sub_move).squares) == 2
            for sub_move in self.get_moves(check_only_check=True))

    def get_moves(self, check_only_check=False):
        for square in SQUARES:
            if self[square] is None or self[square].colour != self.move:
                continue
            for move in self[square].get_moves(self,
                                               captures_only=check_only_check):
                move = (square, move)
                if not check_only_check:
                    new_pos = self.get_moved(move)
                    if not new_pos.is_safe():
                        break
                    else:
                        yield move
                else:
                    yield move
        # HACK
        if check_only_check:
            return
        for way in self.castling:
            if way.isupper() != (self.move == "w"):
                continue
            bad_squares = []
            for square in SQUARES:
                if isinstance(self[square],
                              King) and self[square].colour == self.move:
                    bad_squares.append(square)
                    break
            for i in range(1, 5):
                try:
                    new_square = bad_squares[0] + (
                        i if way.casefold() == "k" else -i,
                        0,
                    )
                except ValueError:
                    break
                bad_squares.append(new_square)
                if isinstance(self[new_square], Rook):
                    break
            if all((self[s] is None or s in (bad_squares[0], bad_squares[-1]))
                   and self.get_moved((bad_squares[0], s)).is_safe()
                   for s in bad_squares):
                yield (bad_squares[0], bad_squares[2])

    def get_moved(self, move):
        board = Board()
        board.squares = []
        for square in SQUARES:
            if move[1] == square:
                board.squares.append(self[move[0]].get_moved(move[1]))
            elif move[0] == square:
                board.squares.append(None)
            else:
                board.squares.append(self[square])
        delta = move[1] - move[0]
        if delta in {(2, 0), (-2, 0)} and isinstance(self[move[0]], King):
            for square in range(3):
                square = move[1] + (square * delta[0] // 2, 0)
                if isinstance(self[square],
                              Rook) and self[square].colour == self.move:
                    if isinstance(board[square], Rook):
                        board[square] = None
                    move_to = move[1] - (delta[0] // 2, 0)
                    board[move_to] = self[square].get_moved(move_to)
                    break
        board.move = "b" if self.move == "w" else "w"
        board.castling = []
        for way in "KQkq":
            if way not in self.castling:
                continue
            their_move = way.isupper() == (self.move == "w")
            if their_move and (isinstance(self[move[0]], King) or
                               (isinstance(self[move[0]], Rook) and
                                ((move[0 if their_move else 1].file_num < 4) ==
                                 (way.upper() == "K")))):
                continue
            board.castling.append(way)

        if self[move[1]] is not None or isinstance(self[move[0]], Pawn):
            board.half_moves = 0
        else:
            board.half_moves = self.half_moves + 1
        board.full_moves = self.full_moves + (1 if self.move == "b" else 0)
        if isinstance(self[move[0]], Pawn) and abs(
            (move[1] - move[0])[1]) == 2:
            board.en_passant = str(move[0] + (0, (move[1] - move[0])[1] // 2))
        else:
            board.en_passant = "-"
        return board


def initial_states():
    states = set((a, y) for y in (True, False) for x in "QMC"
                 for a in itertools.permutations("RRBBNN" + x))
    for state, flip in states:
        if (state[3] == "R" or state[4] == "R"
                or state[:4].count("R") != state[4:].count("R")):
            continue
        state = f"{''.join(state[:4])}K{''.join(state[4:])}"
        if state[state.index("B") + 1:].index("B") % 2:
            continue
        pieces = [
            state,
            "P" * 8,
            " " * 8,
            " " * 8,
            " " * 8,
            " " * 8,
            "p" * 8,
            (state if flip else state[::-1]).lower(),
        ]
        board = Board()
        for square in SQUARES:
            board[square] = piece_from_letter(
                pieces[square.rank_num][square.file_num], square)
        yield board


STANDARD = Board()
pieces = [
    "RNBQKBNR",
    "P" * 8,
    " " * 8,
    " " * 8,
    " " * 8,
    " " * 8,
    "p" * 8,
    "rnbqkbnr",
]
for square in SQUARES:
    STANDARD[square] = piece_from_letter(
        pieces[square.rank_num][square.file_num], square)


def get_fens(state, depth=1):
    if depth == 0:
        return {state.to_fen()}
    fens = set()
    for sub_state in state.get_moves():
        fens |= get_fens(state.get_moved(sub_state), depth - 1)
    return fens


def main():
    states = list(initial_states())
    for state in tqdm(states):
        fens = get_fens(state, 4)
        with open(
                f"positions/{state.to_fen().replace('/', '_').split(' ')[0]}.fen",
                "w") as f:
            f.write("\n".join(fens))


if __name__ == "__main__":
    main()