import numpy as np
import torch
import matplotlib.pyplot as plt
import copy
import tqdm
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
import chess

# puzzle presets
board = chess.Board()
files = ["./ww.pgn", "./bw.pgn", "./drew.pgn"]


def board_data(board):
    board_array = np.zeros((8, 8, 13), dtype=np.int8)
    for i in range(64):
        piece = board.piece_at(i)
        if piece is not None:
            color = int(piece.color)
            piece_type = piece.piece_type - 1
            board_array[i // 8][i % 8][piece_type + 6 * color] = 1
        else:
            board_array[i // 8][i % 8][-1] = 1
    board_array = board_array.flatten()
    # board_array.shape = 832 (1D)
    return board_array




# load the numpy arrays (if any), these arrays are generated once and saved locally so that time is not wasted reloading every single time
# the files are only there because during development of the code a lot of the time will be lost through loading games afresh
# getting the number of games needed to analyse
# set up training data
train_data = np.empty(
    (0, 832)
)  # 832 is the size of the flattened board in one hot encoding
train_target = []
for file in files:
    with open(file, "r") as f:
        contents = f.read()
    contents = contents.split(" \n")
    for game in tqdm.tqdm(contents, desc="each game"):
        game = game.split(" \n")
        g = "".join(game)
        game = g.split(" ")
        game = [x for x in game if x]  # remove empty strings
        MAX_MOMENTS = min(20, len(game))
        unsampled_idx = [x for x in range(1, len(game))]
        for move in range(MAX_MOMENTS - 1):
            board = chess.Board()
            up_to = np.random.choice(unsampled_idx)
            unsampled_idx.remove(up_to)
            moment = game[:up_to]
            counter = 0
            for move in moment:
                board.push(chess.Move.from_uci(move))
                counter += 1
            matrix_game = np.array([board_data(board)])  # game after one hot encoding
            train_data = np.append(train_data, matrix_game, axis=0)
            status = 0
            if (counter % 2 == 0 and file == "./ww.pgn") or (
                counter % 2 == 1 and file == "./bw.pgn"
            ):
                status = 1
            elif (counter % 2 == 0 and file == "./bw.pgn") or (
                counter % 2 == 1 and file == "./ww.pgn"
            ):
                status = -1
            elif file == "./drew.pgn":
                status = 0
            train_target.append(status)
train_target = np.array(train_target, dtype=np.float32).reshape(-1, 1)