import chess.pgn
import os
from chess import *
from my_graph import *

path = 'pgns/'

positions = {'a1': A1, 'a2': A2, 'a3': A3, 'a4': A4, 'a5': A5, 'a6': A6, 'a7': A7, 'a8': A8,
             'b1': B1, 'b2': B2, 'b3': B3, 'b4': B4, 'b5': B5, 'b6': B6, 'b7': B7, 'b8': B8,
             'c1': C1, 'c2': C2, 'c3': C3, 'c4': C4, 'c5': C5, 'c6': C6, 'c7': C7, 'c8': C8,
             'd1': D1, 'd2': D2, 'd3': D3, 'd4': D4, 'd5': D5, 'd6': D6, 'd7': D7, 'd8': D8,
             'e1': E1, 'e2': E2, 'e3': E3, 'e4': E4, 'e5': E5, 'e6': E6, 'e7': E7, 'e8': E8,
             'f1': F1, 'f2': F2, 'f3': F3, 'f4': F4, 'f5': F5, 'f6': F6, 'f7': F7, 'f8': F8,
             'g1': G1, 'g2': G2, 'g3': G3, 'g4': G4, 'g5': G5, 'g6': G6, 'g7': G7, 'g8': G8,
             'h1': H1, 'h2': H2, 'h3': H3, 'h4': H4, 'h5': H5, 'h6': H6, 'h7': H7, 'h8': H8
             }

class LoadPgnToGraph:
    def __init__(self, file_name=None, directory=None, number_of_games=None):
        self.directory = directory
        if file_name:
            self.file_name = file_name
            self.player_name = ''.join([i for i in file_name[:-4] if i not in "0123456789"])
            g = MyGraph(file_name[:-4])
        else:
            name = directory.split('/')[-1]
            g = MyGraph(name)
        self.number_of_games = number_of_games if number_of_games != None else self.number_of_games_in_file()
        self.g = g

    def number_of_games_in_file(self):
        pgn = open(path + self.file_name)
        game = chess.pgn.read_game(pgn)
        number_of_games = 0
        while game:
            game = chess.pgn.read_game(pgn)
            number_of_games += 1
        return number_of_games

    def load_players_games(self):
        print("\nProcessing " + self.file_name + '...')
        load_bar = '_' * 10
        print(load_bar)

        game_number = 0
        bar_counter = 1
        is_player_white = False
        pgn = open(path + self.file_name)
        game = chess.pgn.read_game(pgn)

        while game:

            if game_number == self.number_of_games:
                break

            if self.player_name in game.headers["Black"]:
                is_player_white = False
            else:
                is_player_white = True

            board = game.board()

            previous_vertex = None

            for move in game.mainline_moves():
                position1 = str(move)[:2]
                position2 = str(move)[2:]
                piece = str(board.piece_at(positions[position1]))
                board.push(move)

                if piece == piece.lower() and is_player_white:
                    continue
                if piece == piece.upper() and not is_player_white:
                    continue

                vertex = self.g.add_vertex(piece, position1, position2)
                vertex_id = self.g.vertex_id(vertex)
                if previous_vertex:
                    previous_vertex_id = self.g.vertex_id(previous_vertex)
                    self.g.add_edge(previous_vertex, vertex_id)
                previous_vertex = vertex

            if self.number_of_games // bar_counter == 10:
                print('#', end="")
                bar_counter = 0

            game = chess.pgn.read_game(pgn)
            bar_counter += 1
            game_number += 1

        print("\nFinished" +  self.file_name + '.\n')

    def load_games(self):
        game_number = 0
        for file_name in os.listdir(self.directory):
            if not file_name.endswith(".pgn"):
                continue

            print('Processing ' + self.directory + '/' + file_name + '...')
            load_bar = '_'*10
            print(load_bar)

            bar_counter = 1
            is_player_white = False
            pgn = open(self.directory + '/' + file_name)
            game = chess.pgn.read_game(pgn)

            while game:

                if game_number == self.number_of_games:
                    return
                    # break

                board = game.board()

                previous_vertex = None

                for move in game.mainline_moves():
                    position1 = str(move)[:2]
                    position2 = str(move)[2:]
                    piece = str(board.piece_at(positions[position1]))
                    board.push(move)

                    vertex = self.g.add_vertex(piece, position1, position2)
                    vertex_id = self.g.vertex_id(vertex)
                    if previous_vertex:
                        previous_vertex_id = self.g.vertex_id(previous_vertex)
                        self.g.add_edge(previous_vertex, vertex_id)
                    previous_vertex = vertex

                if self.number_of_games // bar_counter == 10:
                    print('#', end="")
                    bar_counter = 0


                game = chess.pgn.read_game(pgn)
                bar_counter += 1
                game_number += 1

            print('\nFinished ' + self.directory + '/' + file_name + '.\n')

    def count_CC(self, number_of_games):
        game_number = 0
        self.g = MyGraph('cc')

        for file_name in os.listdir(self.directory):
            if not file_name.endswith(".pgn"):
                continue

            print('Processing ' + self.directory + '/' + file_name + '...')
            load_bar = '_' * 10
            print(load_bar)


            bar_counter = 1
            pgn = open(self.directory + '/' + file_name)
            game = chess.pgn.read_game(pgn)

            while game:

                if game_number == number_of_games:
                    return self.g.cluster_coeficient()

                board = game.board()

                previous_vertex = None

                for move in game.mainline_moves():
                    position1 = str(move)[:2]
                    position2 = str(move)[2:]
                    piece = str(board.piece_at(positions[position1]))
                    board.push(move)

                    vertex = self.g.add_vertex(piece, position1, position2)
                    vertex_id = self.g.vertex_id(vertex)
                    if previous_vertex:
                        previous_vertex_id = self.g.vertex_id(previous_vertex)
                        self.g.add_edge(previous_vertex, vertex_id)
                    previous_vertex = vertex

                if self.number_of_games // bar_counter == 10:
                    print('#', end="")
                    bar_counter = 0

                game = chess.pgn.read_game(pgn)
                bar_counter += 1
                game_number += 1

            print('\nFinished ' + self.directory + '/' + file_name + '.\n')

    def map_ntice(self, n):
        map_moves = dict()
        pole = [None] * n

        for file_name in os.listdir(self.directory):

            pgn = open(self.directory + '/' + file_name)
            game = chess.pgn.read_game(pgn)

            while game:
                board = game.board()

                for move in game.mainline_moves():
                    for i in range(len(pole) - 1):
                        pole[i] = pole[i + 1]

                    position1 = str(move)[:2]
                    position2 = str(move)[2:]
                    try:
                        piece = str(board.piece_at(positions[position1]))
                    except:
                        continue
                    pole[-1] = piece + position1 + "-" + piece + position2

                    if None not in pole:
                        full_name = ' --- '.join(pole)

                        if (full_name not in map_moves):
                            map_moves[full_name] = 1
                        else:
                            map_moves[full_name] += 1

                    board.push(move)

                game = chess.pgn.read_game(pgn)

        out = sorted(map_moves.items(), key=lambda x: x[1])
        return out

    def visualize(self, layout_type, save=False):
        self.g.visualize(layout_type, save)

    def export(self, name=None):
        self.g.export_to_gml(name)
        self.g.export_to_xlsx(name)