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

# puzzle presets
MAX_MOMENTS = 10
TOTAL_GAMES = 4
board = chess.Board()
files = ["./a.txt","./ab.txt", "./ac.txt"]

# pre-training loop

final_df = pd.DataFrame()

# set up training data

for file in files:
    with open(file, "r") as f:
        contents = f.read()

    contents = contents.split(" \n")
    df_add = pd.DataFrame(columns=["moves", "status"])

    for game in contents:
        if file == "./a.txt":  # NOTE: change the filename
            d = {"moves": game, "status": "won"}
            df_add.loc[len(df_add)] = d
        elif file == "./ab.txt":
            d = {"moves": game, "status": "lost"}
            df_add.loc[len(df_add)] = d
        elif file == "./ac.txt":
            d = {"moves": game, "status": "drew"}
            df_add.loc[len(df_add)] = d
    final_df = pd.concat([final_df, df_add], ignore_index=True)
# define function that takes chess board and turns into AI-understandable matrix


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()
    return board_array


game = 0
train_df = pd.DataFrame(
    columns=[
        "board_data",
        "status",
    ]
)

for index, row in final_df.iterrows():
    moves = row["moves"].split(" ")
    status = row["status"]
    moves = [x for x in moves if x]  # removes empty strings in list
    if len(moves) <= MAX_MOMENTS:
        MAX_MOMENTS = len(moves)
    unsampled_idx = [x for x in range(len(moves))]
    unsampled_idx.pop(0)
    for _ in range(MAX_MOMENTS - 1):
        board = chess.Board()
        up_to = np.random.choice(unsampled_idx)
        unsampled_idx.remove(up_to)
        moment = moves[:up_to]
        df_add = pd.DataFrame(
            columns=[
                "board_data",
                "status",
            ]
        )
        for move in moment:
            board.push(chess.Move.from_uci(move))
        ai_moves = board_data(board)

        counter = 0
        d = {
            "board_data": ai_moves,
            "status": status,
        }
        df_add.loc[len(df_add)] = d
        train_df = pd.concat([train_df, df_add], ignore_index=True)
    game += 1

# preprocessing data

train_df["status"] = train_df["status"].map(
    {"won": 1.0, "lost": -1.0, "drew": 0.0}
)  # target
X = np.array([x for x in train_df["board_data"]])
y = np.array([x for x in train_df["status"]])
# train-test split for model evaluation
X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.7, shuffle=True)
# Convert to 2D PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)


class EvalNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(832, 832)
        self.hidden_layers = nn.ModuleList()
        for _ in range(100):
            self.hidden_layers.append(nn.Linear(832, 832))
        self.fc_out = nn.Tanh()
        self.optimizer = optim.AdamW(self.parameters(), lr=1e-3)
        self.scheduler = optim.lr_scheduler.StepLR(
            self.optimizer, step_size=10, gamma=0.5
        )
        self.loss = nn.MSELoss()

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        for layer in self.hidden_layers:
            x = torch.relu(layer(x))
        x = self.fc_out(x)
        return x

    def load_weights(self, path):
        self.load_state_dict(torch.load(path))


model = EvalNN()
train_losses = []
test_losses = []
num_epochs = 5

for epoch in range(num_epochs):
    # Training
    model.train()
    train_loss = 0.0
    for i in tqdm.tqdm(range(X_train.shape[0])):
        inputs = X_train[i]
        targets = y_train[i]
        model.optimizer.zero_grad()
        outputs = model(inputs)
        loss = model.loss(outputs, targets)
        loss.backward()
        model.optimizer.step()
        train_loss += loss.item()
    train_loss /= X_train.shape[0]
    train_losses.append(train_loss)

    # Validation
    model.eval()
    test_loss = 0.0
    with torch.no_grad():
        for i in range(X_test.shape[0]):
            inputs = X_test[i]
            targets = y_test[i]
            outputs = model(inputs)
            loss = model.loss(outputs, targets)
            test_loss += loss.item()
        test_loss /= X_test.shape[0]
        test_losses.append(test_loss)

    # Learning rate scheduling
    model.scheduler.step()

    print(
        f"Epoch {epoch+1}/{num_epochs}: Train Loss = {train_loss:.6f}, Test Loss = {test_loss:.6f}"
    )
plt.plot(train_losses, label="Train Loss")
plt.plot(test_losses, label="Test Loss")
plt.legend()
plt.show()
plt.savefig('ai-eval-losses.png')